From de7bc61178d418cc3aa2a8c75c136a1a45b4708a Mon Sep 17 00:00:00 2001
From: "TOKISAKIX\\21168" <2116884726@qq.com>
Date: Thu, 12 Sep 2024 10:17:56 +0800
Subject: [PATCH] add optimizer

---
 README.md               |  2 +-
 README_cn.md            |  2 +-
 project/tensor/model.py |  2 +-
 project/tensor/run.py   |  4 ++--
 terox/optim/__init__.py |  1 +
 terox/optim/adam.py     | 31 +++++++++++++++++++++++++++++++
 terox/optim/optim.py    | 17 +++++++++++++++++
 terox/optim/sgd.py      | 20 ++++++++++++++++++++
 8 files changed, 74 insertions(+), 5 deletions(-)
 create mode 100644 terox/optim/__init__.py
 create mode 100644 terox/optim/adam.py
 create mode 100644 terox/optim/optim.py
 create mode 100644 terox/optim/sgd.py

diff --git a/README.md b/README.md
index b2f3e37..4ebb7fb 100644
--- a/README.md
+++ b/README.md
@@ -12,7 +12,7 @@ Terox is a tiny Python package that provides some features:
 - [x] Control the parameters and the model.
 - [ ] Provides common computing functions for deep learning.
 - [ ] Provides common deep learning components.
-- [ ] Provides deep learning model optimizer.
+- [x] Provides deep learning model optimizer.
 - [ ] Accelerate computing on CPU and GPU.
 - [ ] Support distributed computing.
 
diff --git a/README_cn.md b/README_cn.md
index 5791296..0f0f519 100644
--- a/README_cn.md
+++ b/README_cn.md
@@ -12,7 +12,7 @@ Terox 是一个很精简的 Python 包，它提供了一些特性:
 - [x] 便捷控制参数和模型。
 - [ ] 提供深度学习常用的计算函数。
 - [ ] 提供常用的深度学习组件。
-- [ ] 提供深度学习模型优化器。
+- [x] 提供深度学习模型优化器。
 - [ ] 提高在 CPU 和 GPU 上的计算速度。
 - [ ] 支持分布式计算。
 
diff --git a/project/tensor/model.py b/project/tensor/model.py
index fdd52be..1ce6e04 100644
--- a/project/tensor/model.py
+++ b/project/tensor/model.py
@@ -34,7 +34,7 @@ def forward(self, inputs:Tensor) -> Tensor:
         out = out.sigmoid()
         return out
     
-class SGD():
+class GD():
     def __init__(self, parameters:List[Parameter], lr:float) -> None:
         self.parameters = parameters
         self.lr = lr
diff --git a/project/tensor/run.py b/project/tensor/run.py
index 7ba824b..8d2a864 100644
--- a/project/tensor/run.py
+++ b/project/tensor/run.py
@@ -4,7 +4,7 @@
 from terox.tensor import Tensor
 from terox.module import Module
 
-from model import ScalarIrisClassifyModel, SGD
+from model import ScalarIrisClassifyModel, GD
 from dataset import getDataSet
 from function import MSELoss
 
@@ -25,7 +25,7 @@ def test(model:Module, dataset:List) -> float:
     dataset   = getDataSet(N)
     model     = ScalarIrisClassifyModel(in_feature=2, hidden_feature=128, out_feature=1)
     criterion = MSELoss
-    optimizer = SGD(model.parmeters(), LR)
+    optimizer = GD(model.parmeters(), LR)
     print(f"[INFO] Start Acc:{test(model, dataset) * 100:.2f}%")
 
     for epoch in range(EPOCHS):
diff --git a/terox/optim/__init__.py b/terox/optim/__init__.py
new file mode 100644
index 0000000..b437bfb
--- /dev/null
+++ b/terox/optim/__init__.py
@@ -0,0 +1 @@
+from .sgd import SGD
\ No newline at end of file
diff --git a/terox/optim/adam.py b/terox/optim/adam.py
new file mode 100644
index 0000000..af44a89
--- /dev/null
+++ b/terox/optim/adam.py
@@ -0,0 +1,31 @@
+import math
+from typing import List
+
+from terox.module import Parameter
+from .optim import Optimizer
+
+class Adam(Optimizer):
+    def __init__(self, parameters: List[Parameter], lr: float, beta1: float = 0.9, beta2: float = 0.999, epsilon: float = 1e-8) -> None:
+        super().__init__(parameters, lr)
+        self.beta1 = beta1
+        self.beta2 = beta2
+        self.epsilon = epsilon
+        self.m = {id(param): 0.0 for param in parameters}
+        self.v = {id(param): 0.0 for param in parameters}
+        self.t = 0
+        return
+
+    def step(self) -> None:
+        self.t += 1
+        for parameter in self.parameters:
+            grad = parameter.value()._gradient._item
+            param_id = id(parameter)
+            
+            self.m[param_id] = self.beta1 * self.m[param_id] + (1 - self.beta1) * grad
+            self.v[param_id] = self.beta2 * self.v[param_id] + (1 - self.beta2) * (grad ** 2)
+            
+            m_hat = self.m[param_id] / (1 - self.beta1 ** self.t)
+            v_hat = self.v[param_id] / (1 - self.beta2 ** self.t)
+            
+            parameter.value()._item -= self.lr * m_hat / (math.sqrt(v_hat) + self.epsilon)
+        return
\ No newline at end of file
diff --git a/terox/optim/optim.py b/terox/optim/optim.py
new file mode 100644
index 0000000..74cd4ea
--- /dev/null
+++ b/terox/optim/optim.py
@@ -0,0 +1,17 @@
+from typing import List
+
+from terox.module import Parameter
+
+class Optimizer():
+    def __init__(self, parameters: List[Parameter], lr: float) -> None:
+        self.parameters = parameters
+        self.lr = lr
+        return
+
+    def zero_grad(self) -> None:
+        for parameter in self.parameters:
+            parameter.value()._zeroGrad()
+        return
+    
+    def step(self) -> None:
+        raise NotImplementedError
\ No newline at end of file
diff --git a/terox/optim/sgd.py b/terox/optim/sgd.py
new file mode 100644
index 0000000..e7c8c34
--- /dev/null
+++ b/terox/optim/sgd.py
@@ -0,0 +1,20 @@
+from typing import List
+
+from terox.module import Parameter
+from .optim import Optimizer
+    
+class SGD(Optimizer):
+    def __init__(self, parameters: List[Parameter], lr: float, momentum: float = 0.0) -> None:
+        super().__init__(parameters, lr)
+        self.momentum = momentum
+        self.velocity = {id(param): 0.0 for param in parameters}
+        return
+    
+    def step(self) -> None:
+        for parameter in self.parameters:
+            grad = parameter.value()._gradient._item
+            param_id = id(parameter)
+            
+            self.velocity[param_id] = self.momentum * self.velocity[param_id] - self.lr * grad
+            parameter.value()._item += self.velocity[param_id]
+        return
\ No newline at end of file