Initial commit

test.py

+from tqdm import tqdm
+import torch
+from torch import Tensor
+from torch.nn import Module, Linear, MSELoss
+from torch.optim import Adam
+
+
+class Model(Module):
+    def __init__(self):
+        super(Model, self).__init__()
+
+        self.linear = Linear(in_features=1, out_features=1)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.linear(x)
+        return x
+
+
+def main():
+    max_epochs = 100
+    x = torch.randn(1000, 1)
+    y = x * 5.0 - 2.0
+
+    m = Model()
+
+    opt = Adam(m.parameters())
+
+    criterion = MSELoss()
+
+    for _ in tqdm(range(max_epochs)):
+        for i in tqdm(range(len(x)), leave=False):
+            opt.zero_grad()
+
+            loss = criterion(input=m(x[i]), target=y[i])
+            loss.backward()
+
+            opt.step()
+
+    loss = criterion(input=m(x), target=y)
+    print(loss)
+    print(m.linear.weight)
+    print(m.linear.bias)
+
+
+if __name__ == '__main__':
+    main()