src/machine_learning/neural_network/learning.py

@@ -3,7 +3,6 @@ import torch.nn as nn
 from torch.utils.data import DataLoader, random_split
 
 from net import NeuralNetwork
-from database import create_training_data
 from Dataset import ImageDataset
 import os
 
@@ -18,19 +17,13 @@ annonation_file = os.path.abspath('annotations.csv')
 
 dataset = ImageDataset(annonation_file, img_dir)
 
-print(len(dataset))
 trainset, testset = random_split(dataset, [1031, 200])
 
 batch_size = 64
 
 # Create data loaders.
-train_dataloader = DataLoader(trainset, batch_size=batch_size, shuffle=True)
-test_dataloader = DataLoader(testset, batch_size=batch_size, shuffle=True)
-
-for X, y in test_dataloader:
-    print("Shape of X [N, C, H, W]: ", X.shape)
-    print("Shape of y: ", y.shape, y.dtype)
-    break
+train_dataloader = DataLoader(trainset, shuffle=True)
+test_dataloader = DataLoader(testset, shuffle=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print("Using {} device".format(device))
@@ -48,7 +41,7 @@ def train(dataloader, model, loss_fn, optimizer):
         X, y = X.to(device), y.to(device)
 
         # Compute prediction error
-        pred = model(X)
+        pred = model(X.float())
         loss = loss_fn(pred, y)
 
         # Backpropagation
@@ -68,7 +61,7 @@ def test(dataloader, model):
     with torch.no_grad():
         for X, y in dataloader:
             X, y = X.to(device), y.to(device)
-            pred = model(X)
+            pred = model(X.float())
             test_loss += loss_fn(pred, y).item()
             correct += (pred.argmax(1) == y).type(torch.float).sum().item()
     test_loss /= size

src/machine_learning/neural_network/net.py

@@ -6,7 +6,7 @@ class NeuralNetwork(nn.Module):
         super(NeuralNetwork, self).__init__()
         self.flatten = nn.Flatten()
         self.linear_relu_stack = nn.Sequential(
-            nn.Linear(28*28, 512),
+            nn.Linear(28*28*3, 512),
             nn.ReLU(),
             nn.Linear(512, 512),
             nn.ReLU(),