[fix] Validation step in pytorch_track example

examples/pytorch_track.py

@@ -10,16 +10,18 @@
 
 
 # Initialize a new Run
-aim_run = Run()
+aim_run = Run(capture_terminal_logs=True)
+
+logging.getLogger()
 
 # moving model to gpu if available
 device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
 # Hyper parameters
 num_epochs = 5
 num_classes = 10
-batch_size = 16
-learning_rate = 0.01
+batch_size = 32
+learning_rate = 0.001
 
 # aim - Track hyper parameters
 aim_run['hparams'] = {
@@ -40,27 +42,20 @@
 test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)
 
 
-# Convolutional neural network (two convolutional layers)
+# Convolutional neural network (one convolutional layer)
 class ConvNet(nn.Module):
     def __init__(self, num_classes=10):
         super(ConvNet, self).__init__()
         self.layer1 = nn.Sequential(
-            nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2),
-            nn.BatchNorm2d(16),
+            nn.Conv2d(1, 2, kernel_size=5, stride=1, padding=2),
+            nn.BatchNorm2d(2),
             nn.ReLU(),
             nn.MaxPool2d(kernel_size=2, stride=2),
         )
-        self.layer2 = nn.Sequential(
-            nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2),
-            nn.BatchNorm2d(32),
-            nn.ReLU(),
-            nn.MaxPool2d(kernel_size=2, stride=2),
-        )
-        self.fc = nn.Linear(7 * 7 * 32, num_classes)
+        self.fc = nn.Linear(2 * 14 * 14, num_classes)
 
     def forward(self, x):
         out = self.layer1(x)
-        out = self.layer2(out)
         out = out.reshape(out.size(0), -1)
         out = self.fc(out)
         return out
@@ -88,21 +83,19 @@ def forward(self, x):
         loss.backward()
         optimizer.step()
 
-        if i % 30 == 0:
-            logging.info(
-                'Epoch [{}/{}], Step [{}/{}], ' 'Loss: {:.4f}'.format(
-                    epoch + 1, num_epochs, i + 1, total_step, loss.item()
-                )
+
+        aim_run.log_info(
+            f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Train Loss: {loss.item():.4f}"
             )
 
-            # aim - Track model loss function
-            correct = 0
-            total = 0
-            _, predicted = torch.max(outputs.data, 1)
-            total += labels.size(0)
-            correct += (predicted == labels).sum().item()
-            acc = 100 * correct / total
+        correct = 0
+        total = 0
+        _, predicted = torch.max(outputs.data, 1)
+        total += labels.size(0)
+        correct += (predicted == labels).sum().item()
+        acc = 100 * correct / total
 
+        if i % 30 == 0: 
             # aim - Track metrics
             items = {'accuracy': acc, 'loss': loss}
             aim_run.track(items, epoch=epoch, context={'subset': 'train'})
@@ -111,22 +104,29 @@ def forward(self, x):
             track_params_dists(model, aim_run)
             track_gradients_dists(model, aim_run)
 
-            # TODO: Do actual validation
-            if i % 300 == 0:
-                aim_run.track(loss.item(), name='loss', epoch=epoch, context={'subset': 'val'})
-                aim_run.track(acc, name='accuracy', epoch=epoch, context={'subset': 'val'})
-
-
-# Test the model
-with torch.inference_mode():
-    correct = 0
-    total = 0
-    for images, labels in test_loader:
+    model.eval()
+    for i, (images, labels) in enumerate(test_loader):
         images = images.to(device)
         labels = labels.to(device)
+
+        # Forward pass
         outputs = model(images)
+        loss = criterion(outputs, labels)
+
+        logging.info(
+            f"Epoch [{epoch+1}/{num_epochs}], Step [{i+1}/{total_step}], Valid Loss: {loss.item():.4f}"
+            )
+
+        correct = 0
+        total = 0
         _, predicted = torch.max(outputs.data, 1)
         total += labels.size(0)
         correct += (predicted == labels).sum().item()
+        acc = 100 * correct / total
+
+        if i % 30 == 0:
+            # aim - Track metrics
+            items = {'accuracy': acc, 'loss': loss}
+            aim_run.track(items, epoch=epoch, context={'subset': 'test'})
 
-    logging.info('Test Accuracy: {} %'.format(100 * correct / total))
+logging.info('Training finished')