small fixes

This commit is contained in:
MarRac 2024-05-25 18:41:25 +02:00
parent b5d25d710d
commit c363b09f85
3 changed files with 13 additions and 5 deletions

Binary file not shown.

View File

@ -1,4 +1,6 @@
import torch.nn as nn
import torch
import torch.nn.functional as F
class Neural_Network_Model(nn.Module):
@ -16,5 +18,4 @@ class Neural_Network_Model(nn.Module):
x = self.fc2(x)
x = torch.relu(x)
x = self.out(x)
F.log_softmax(x, dim=-1)
return x
return F.log_softmax(x, dim=-1)

View File

@ -1,15 +1,16 @@
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision import datasets, transforms, utils
from torchvision.transforms import Compose, Lambda, ToTensor
import matplotlib.pyplot as plt
import numpy as np
from model import *
device = torch.device('cuda')
#data transform to tensors:
data_transformer = transforms.Compose
([
data_transformer = transforms.Compose([
transforms.Resize((150, 150)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
@ -46,6 +47,9 @@ def train(model, dataset, iter=100, batch_size=64):
loss = criterion(output, labels.to(device))
loss.backward()
optimizer.step()
if epoch % 10 == 0:
print('epoch: %3d loss: %.4f' % (epoch, loss))
#function for getting accuracy
def accuracy(model, dataset):
model.eval()
@ -57,6 +61,9 @@ def accuracy(model, dataset):
return correct.float() / len(dataset)
model = Neural_Network_Model()
model.to(device)
train(model, train_set)
print(accuracy(model, test_set))