import glob from src.torchvision_resize_dataset import combined_dataset, images_path, classes import src.data_model from torch.optim import Adam import torch import torch.nn as nn from torch.utils.data import DataLoader device = torch.device("cuda" if torch.cuda.is_available() else "cpu") train_loader = DataLoader( combined_dataset, #dataset of images batch_size=256, # accuracy shuffle=True # rand order ) 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(images_path, '*.jpg')) 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")