Dodano funkcję train i accuracy oraz model z dwoma ukrytymi warswtami

source/NN/model.py

@@ -0,0 +1,20 @@
+import torch.nn as nn
+
+
+class Neural_Network_Model(nn.Module):
+    def __init__(self, num_classes=5,hidden_layer1 = 100,hidden_layer2 = 100):
+        super(Neural_Network_Model, self).__init__()
+        self.fc1 = nn.Linear(150*150*3,hidden_layer1)
+        self.fc2 = nn.Linear(hidden_layer1,hidden_layer2)
+        self.out = nn.Linear(hidden_layer2,num_classes)
+# two hidden layers
+
+    def forward(self, x):
+        x = x.view(-1, 150*150*3)
+        x = self.fc1(x)
+        x = torch.relu(x)
+        x = self.fc2(x)
+        x = torch.relu(x)
+        x = self.out(x)
+        F.log_softmax(x, dim=-1)
+        return x

source/NN/neural_network.py

@@ -5,7 +5,7 @@ from torchvision import datasets, transforms
 from torchvision.transforms import Compose, Lambda, ToTensor
 import matplotlib.pyplot as plt
 import numpy as np
-
+from model import *
 
 #data transform to tensors:
 data_transformer = transforms.Compose
@@ -17,18 +17,46 @@ data_transformer = transforms.Compose
 
 
 #loading data:
-train_set = datasets.ImageFolder(root='source/resources/train', transform=data_transformer)
-test_set = datasets.ImageFolder(root='source/resources/test', transform=data_transformer)
+train_set = datasets.ImageFolder(root='resources/train', transform=data_transformer)
+test_set = datasets.ImageFolder(root='resources/test', transform=data_transformer)
 
 
 #to mozna nawet przerzucic do funkcji train:
-train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)
-test_loader = DataLoader(test_set, batch_size=32, shuffle=True, num_workers=2)
+#train_loader = DataLoader(train_set, batch_size=32, shuffle=True, num_workers=2)
+#test_loader = DataLoader(test_set, batch_size=32, shuffle=True, num_workers=2)
 
 #test if classes work properly:
-print(train_set.classes)  
-print(train_set.class_to_idx)  
-print(train_set.targets[3002])
+#print(train_set.classes)
+#print(train_set.class_to_idx)
+#print(train_set.targets[3002])
+
 
 #loading your own image: <-- zrobię to na koniec - wrzucanie konkretnego obrazka aby uzyskac wynik
+#function for training model
+def train(model, dataset, iter=100, batch_size=64):
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
+    criterion = nn.NLLLoss()
+    train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+    model.train()
 
+    for epoch in range(iter):
+        for inputs, labels in train_loader:
+            optimizer.zero_grad()
+            output = model(inputs.to(device))
+            loss = criterion(output, labels.to(device))
+            loss.backward()
+            optimizer.step()
+#function for getting accuracy
+def accuracy(model, dataset):
+    model.eval()
+    correct = sum([
+        (model(inputs.to(device)).argmax(dim=1) == labels.to(device)).sum()
+        for inputs, labels in DataLoader(dataset, batch_size=64, shuffle=True)
+    ])
+
+    return correct.float() / len(dataset)
+
+
+model = Neural_Network_Model()
+train(model, train_set)
+print(accuracy(model, test_set))