feature: Added manual training and PyTorch Lightning training loops

src/{main.py → train.py}

@@ -1,3 +1,5 @@
+import torch
+import numpy as np
 from torch import nn, optim
 from torch.utils.data import random_split
 import pytorch_lightning as pl
@@ -22,13 +24,32 @@ def main():
     validate_dataloader = DataLoader(validate_data, num_workers=2)
     test_dataloader = DataLoader(test_data, num_workers=8)  # My CPU has 8 cores
 
-    # grayscale channels = 1, mnist num_labels = 10
-    net = CNN(input_channels=1, num_classes=10)
+    net = CNN(input_channels=1, num_classes=10).to("cuda")
+    opt = optim.Adam(net.parameters(), lr=1e-4)
+    loss_fn = nn.CrossEntropyLoss()
+    max_epochs = 10
+    for i in range(max_epochs):
+        for idx, batch in enumerate(train_dataloader):
+            x, y = batch
+            x = x.to("cuda")
+            y = y.to("cuda")
+
+            y_pred = net(x).reshape(1, -1)
+            loss = loss_fn(y_pred, y)
+
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
 
-    pl_net = LitTrainer(net, nn.CrossEntropyLoss(), optim.Adam(net.parameters()))
-    trainer = pl.Trainer(limit_train_batches=100, max_epochs=1, default_root_dir="../checkpoints")
+            if idx % 1000 == 0:
+                print(f"Loss: {loss.item()} ({idx} / {len(train_dataloader)})")
 
-    trainer.fit(model=pl_net, train_dataloaders=train_dataloader, val_dataloaders=validate_dataloader)
+    torch.save(net, "../checkpoints/pytorch/version_1.pt")
+
+    # grayscale channels = 1, mnist num_labels = 10
+    trainer = pl.Trainer(limit_train_batches=100, max_epochs=10, default_root_dir="../checkpoints")
+    pl_net = LitTrainer(CNN(input_channels=1, num_classes=10))
+    trainer.fit(pl_net, train_dataloader, validate_dataloader)
     trainer.test(model=pl_net, dataloaders=test_dataloader)