Optimizer torch optimizer performance (#482)

* add torch optimizers

* addressing comments

---------

Co-authored-by: Haifeng Jin <[email protected]>

Author: haifeng-jin

benchmarks/torch_ctl_benchmark/README.md

@@ -1,9 +1,10 @@
 # Benchmark the performance of torch custom training loop
 
-This directory contains benchmarks to compare the performance between Keras and
-Torch while using Torch custom training loop. The benchmark purpose is to
-understand the performance diff resulting from the modeling API choice (Keras
-or Torch).
+This directory contains benchmarks to compare the performance of a Keras model
+and a equivalent Torch model while using the same Torch custom training loop.
+
+The benchmark purpose is to understand the performance diff resulting from the
+modeling API choice (Keras or Torch).
 
 To run the benchmark, use the command below and change to your target:
 

keras_core/backend/torch/optimizers/__init__.py

@@ -0,0 +1 @@
+from keras_core.backend.torch.optimizers.torch_optimizer import TorchOptimizer

keras_core/backend/torch/optimizers/torch_optimizer.py

@@ -0,0 +1,24 @@
+import torch
+
+from keras_core.optimizers.base_optimizer import BaseOptimizer
+
+
+class TorchOptimizer(BaseOptimizer):
+    def __new__(cls, *args, **kwargs):
+        # Import locally to avoid circular imports.
+        from keras_core import optimizers
+        from keras_core.backend.torch.optimizers import torch_sgd
+
+        OPTIMIZERS = {optimizers.SGD: torch_sgd.SGD}
+        if cls in OPTIMIZERS:
+            return OPTIMIZERS[cls](*args, **kwargs)
+        return super().__new__(cls)
+
+    def _apply_weight_decay(self, variables):
+        if self.weight_decay is None:
+            return
+
+        torch._foreach_mul_(
+            [v.value for v in variables if self._use_weight_decay(v)],
+            1 - self.weight_decay * self._get_current_learning_rate(),
+        )

keras_core/backend/torch/optimizers/torch_sgd.py

@@ -0,0 +1,43 @@
+import torch
+
+from keras_core import optimizers
+
+
+class SGD(optimizers.SGD):
+    def _internal_apply_gradients(self, grads_and_vars):
+        grads, trainable_variables = zip(*grads_and_vars)
+
+        self._parallel_update_step(
+            grads,
+            [v.value for v in trainable_variables],
+            self._get_current_learning_rate(),
+        )
+        self.iterations.assign(self.iterations + 1)
+
+    def _parallel_update_step(
+        self,
+        grads,
+        variables,
+        learning_rate,
+    ):
+        if self.momentum != 0:
+            bufs = [
+                self.momentums[self._get_variable_index(variable.value)]
+                for variable in variables
+            ]
+
+            for i in range(len(bufs)):
+                if bufs[i] is None:
+                    bufs[i] = torch.clone(grads[i]).detach()
+
+            torch._foreach_mul_(bufs, self.momentum)
+            torch._foreach_add_(bufs, grads, alpha=-learning_rate)
+
+            if self.nesterov:
+                torch._foreach_add_(variables, grads, alpha=-learning_rate)
+                torch._foreach_add_(variables, bufs, alpha=self.momentum)
+            else:
+                torch._foreach_add_(variables, bufs)
+
+        else:
+            torch._foreach_add_(variables, grads, alpha=-learning_rate)

keras_core/optimizers/optimizer.py

@@ -3,9 +3,13 @@
 from keras_core.optimizers import base_optimizer
 
 if backend.backend() == "tensorflow":
-    from keras_core.backend.tensorflow import optimizer as tf_optimizer
+    from keras_core.backend.tensorflow.optimizer import TFOptimizer
 
-    BackendOptimizer = tf_optimizer.TFOptimizer
+    BackendOptimizer = TFOptimizer
+elif backend.backend() == "torch":
+    from keras_core.backend.torch.optimizers import TorchOptimizer
+
+    BackendOptimizer = TorchOptimizer
 else:
     BackendOptimizer = base_optimizer.BaseOptimizer
 

keras_core/optimizers/sgd_test.py

@@ -21,7 +21,7 @@ def test_config(self):
     def test_single_step(self):
         optimizer = SGD(learning_rate=0.5)
         self.assertEqual(len(optimizer.variables), 2)
-        grads = np.array([1.0, 6.0, 7.0, 2.0])
+        grads = ops.array([1.0, 6.0, 7.0, 2.0])
         vars = backend.Variable([1.0, 2.0, 3.0, 4.0])
         optimizer.build([vars])
         optimizer.apply_gradients(zip([grads], [vars]))
@@ -32,7 +32,7 @@ def test_single_step(self):
 
     def test_weight_decay(self):
         grads, var1, var2, var3 = (
-            np.zeros(()),
+            ops.zeros(()),
             backend.Variable(2.0),
             backend.Variable(2.0, name="exclude"),
             backend.Variable(2.0),
@@ -56,8 +56,8 @@ def test_correctness_with_golden(self):
         optimizer = SGD(nesterov=True)
 
         x = backend.Variable(np.ones([10]))
-        grads = np.arange(0.1, 1.1, 0.1)
-        first_grads = np.full((10,), 0.01)
+        grads = ops.arange(0.1, 1.1, 0.1)
+        first_grads = ops.full((10,), 0.01)
 
         # fmt: off
         golden = np.array(