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

This commit is contained in:
s481808 2024-05-25 16:33:34 +02:00
parent 2ec3f1a89a
commit b5d25d710d
2 changed files with 56 additions and 8 deletions

20
source/NN/model.py Normal file
View File

@ -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

View File

@ -5,7 +5,7 @@ from torchvision import datasets, transforms
from torchvision.transforms import Compose, Lambda, ToTensor from torchvision.transforms import Compose, Lambda, ToTensor
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
import numpy as np import numpy as np
from model import *
#data transform to tensors: #data transform to tensors:
data_transformer = transforms.Compose data_transformer = transforms.Compose
@ -17,18 +17,46 @@ data_transformer = transforms.Compose
#loading data: #loading data:
train_set = datasets.ImageFolder(root='source/resources/train', transform=data_transformer) train_set = datasets.ImageFolder(root='resources/train', transform=data_transformer)
test_set = datasets.ImageFolder(root='source/resources/test', transform=data_transformer) test_set = datasets.ImageFolder(root='resources/test', transform=data_transformer)
#to mozna nawet przerzucic do funkcji train: #to mozna nawet przerzucic do funkcji train:
train_loader = DataLoader(train_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_loader = DataLoader(test_set, batch_size=32, shuffle=True, num_workers=2)
#test if classes work properly: #test if classes work properly:
print(train_set.classes) #print(train_set.classes)
print(train_set.class_to_idx) #print(train_set.class_to_idx)
print(train_set.targets[3002]) #print(train_set.targets[3002])
#loading your own image: <-- zrobię to na koniec - wrzucanie konkretnego obrazka aby uzyskac wynik #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))