From 9172aa397fc309e7205f869b37ffecb61773fe4d Mon Sep 17 00:00:00 2001 From: Woj Date: Sun, 4 Jun 2023 12:38:13 +0200 Subject: [PATCH] learning algoritm --- NeuralNetwork/neural_network_learning.py | 51 +++++++++++++++++-- .../src/torchvision_resize_dataset.py | 1 - 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/NeuralNetwork/neural_network_learning.py b/NeuralNetwork/neural_network_learning.py index a0c75e6..9f71d86 100644 --- a/NeuralNetwork/neural_network_learning.py +++ b/NeuralNetwork/neural_network_learning.py @@ -1,10 +1,10 @@ -from src.torchvision_resize_dataset import combined_dataset +import glob +from src.torchvision_resize_dataset import combined_dataset, letters_path, package_path +import src.data_model +from torch.optim import Adam import torch import torch.nn as nn from torch.utils.data import DataLoader -from torchvision import datasets -from torchvision.transforms import Compose, Lambda, ToTensor -import matplotlib.pyplot as plt device = torch.device("cuda" if torch.cuda.is_available() else "cpu") @@ -15,3 +15,46 @@ train_loader = DataLoader( ) classes = ["package", "list"] +model = src.data_model.DataModel(num_objects=2).to(device) + +#optimizer +optimizer = Adam(model.parameters(), lr=0.001, weight_decay=0.0001) +#loss function +criterion = nn.CrossEntropyLoss() + +num_epochs = 10 +train_size = len(glob.glob(letters_path, '*.jpg')) + len(glob.glob(package_path, '*.png')) + +go_to_accuracy = 0.0 +for epoch in range(num_epochs): + #training on dataset + model.train() + train_accuracy = 0.0 + train_loss = 0.0 + for i, (images, labels) in enumerate(train_loader): + if torch.cuda.is_available(): + images = torch.Variable(images.cuda()) + labels = torch.Variable(labels.cuda()) + # clearing the optimizer gradients + optimizer.zero_grad() + + outputs = model(images) # predoction + loss = criterion(outputs, labels) #loss calculation + loss.backward() + optimizer.step() + + train_loss += loss.cpu().data*images.size(0) + _, prediction = torch.max(outputs.data, 1) + + train_accuracy += int(torch.sum(prediction == labels.data)) + + train_accuracy = train_accuracy/train_size + train_loss = train_loss/train_size + + print('Epoch: '+ str(epoch+1) +' Train Loss: '+ str(int(train_loss)) +' Train Accuracy: '+ str(train_accuracy)) + + if train_accuracy > go_to_accuracy: + go_to_accuracy= train_accuracy + torch.save(model.state_dict(), "best_model.pth") + + diff --git a/NeuralNetwork/src/torchvision_resize_dataset.py b/NeuralNetwork/src/torchvision_resize_dataset.py index 7684909..666d072 100644 --- a/NeuralNetwork/src/torchvision_resize_dataset.py +++ b/NeuralNetwork/src/torchvision_resize_dataset.py @@ -1,7 +1,6 @@ import torchvision.transforms as transforms from torchvision.datasets import ImageFolder from torch.utils.data import ConcatDataset -import os # images have to be the same size for the algorithm to work transform = transforms.Compose([