From 51521dd2ac6f6e1ccd04011f70e869e7ab90a47b Mon Sep 17 00:00:00 2001
From: matixezor <matixezor10@interia.pl>
Date: Sun, 30 May 2021 17:26:20 +0200
Subject: [PATCH] =?UTF-8?q?i=C3=B3sz=20pomagam=20piter?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/machine_learning/neural_network/learning.py | 15 ++++-----------
 src/machine_learning/neural_network/net.py      |  2 +-
 2 files changed, 5 insertions(+), 12 deletions(-)

diff --git a/src/machine_learning/neural_network/learning.py b/src/machine_learning/neural_network/learning.py
index 4e21d0c..a2cc2f7 100644
--- a/src/machine_learning/neural_network/learning.py
+++ b/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
diff --git a/src/machine_learning/neural_network/net.py b/src/machine_learning/neural_network/net.py
index cb6f51f..db51007 100644
--- a/src/machine_learning/neural_network/net.py
+++ b/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(),