# Import necessary packages.import torch
import torch.nn as nn
from torchvision import datasets, models, transforms
from torchsummary import summary
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import time
import os
# Load the Subset-datasets of Caltech 256.
dataset ='./Datasets/'# Set the dataset directory.
train_directory = os.path.join(dataset,'train')
valid_directory = os.path.join(dataset,'valid')
test_directory = os.path.join(dataset,'test')# Number of classes.# num_classes = len(os.listdir(valid_directory))# print(num_classes) # num_classes is 10.# Load Data from folders.
data ={'train':datasets.ImageFolder(root=train_directory,
transform=image_transforms['train']),'valid':datasets.ImageFolder(root=valid_directory,
transform=image_transforms['valid']),'test': datasets.ImageFolder(root=test_directory,
transform=image_transforms['test'])}
# Freeze model parameters, make a preparation for fine-tuning.for param in resnet50.parameters():
param.requires_grad =False
# Change the final layer of ResNet50 for Transfer Learning.
fc_inputs = resnet50.fc.in_features
resnet50.fc = nn.Sequential(nn.Linear(fc_inputs,256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256,len(os.listdir(valid_directory))),# Related to the 10 outputs classes.
nn.LogSoftmax(dim=1))# For using NLLLoss() # Convert model to be used on GPU.
resnet50 = resnet50.to(device)
# Loss and optimizer.
criterion = nn.NLLLoss()
optimizer = torch.optim.Adam(resnet50.parameters())
# Define Train & Validate Utils Function.deftrain_and_validate(model, loss_criterion, optimizer, epochs=25):"""
Function used to train and validate.
Parameters:
:param model: Model to train and validate.
:param loss_criterion: Loss Criterion to minimize.
:param optimizer: Optimizer for computing gradients.
:param epochs: Number of epochs (default=25)
"""
history =[]
best_loss =100000.0
best_epoch =Nonefor epoch inrange(epochs):
epoch_start = time.time()print("Epoch: {}/{}".format(epoch+1, epochs))# Set to training mode.
model.train()# Loss and Accuracy within the epoch.
train_loss =0.0
train_acc =0.0
valid_loss =0.0
valid_acc =0.0for _,(inputs, labels)inenumerate(train_data_loader):# Move the data pair to device.
inputs = inputs.to(device)
labels = labels.to(device)# Forward pass and calculate loss.
outputs = model(inputs)
loss = loss_criterion(outputs, labels)# Backward and optimize.
optimizer.zero_grad()
loss.backward()
optimizer.step()# Compute the total loss for the batch and add it to train_loss.
train_loss += loss.item()* inputs.size(0)# Compute the accuracy.
_, predictions = torch.max(outputs.data,1)
correct_counts = predictions.eq(labels.data.view_as(predictions))# Convert correct_counts to float and then compute the mean.
acc = torch.mean(correct_counts.type(torch.FloatTensor))# Compute total accuracy in the whole batch and add to train_acc.
train_acc += acc.item()* inputs.size(0)#print("Batch number: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}".format(i, loss.item(), acc.item()))# Validation - No gradient tracking needed.with torch.no_grad():# Set to evaluation mode.
model.eval()# Validation loop.for _,(inputs, labels)inenumerate(valid_data_loader):
inputs = inputs.to(device)
labels = labels.to(device)# Forward pass - compute outputs on input data using the model
outputs = model(inputs)# Compute loss
loss = loss_criterion(outputs, labels)# Compute the total loss for the batch and add it to valid_loss
valid_loss += loss.item()* inputs.size(0)# Calculate validation accuracy
_, predictions = torch.max(outputs.data,1)
correct_counts = predictions.eq(labels.data.view_as(predictions))# Convert correct_counts to float and then compute the mean
acc = torch.mean(correct_counts.type(torch.FloatTensor))# Compute total accuracy in the whole batch and add to valid_acc
valid_acc += acc.item()* inputs.size(0)#print("Validation Batch number: {:03d}, Validation: Loss: {:.4f}, Accuracy: {:.4f}".format(j, loss.item(), acc.item()))if valid_loss < best_loss:
best_loss = valid_loss
best_epoch = epoch
# Find average training loss and train accuracy.
avg_train_loss = train_loss/len(data['train'])
avg_train_acc = train_acc/len(data['train'])# Find average training loss and training accuracy.
avg_valid_loss = valid_loss/len(data['valid'])
avg_valid_acc = valid_acc/len(data['valid'])
history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc])
epoch_end = time.time()print("Epoch : {:03d}, Training: Loss - {:.4f}, Accuracy - {:.4f}%, \n\t\tValidation : Loss - {:.4f}, Accuracy - {:.4f}%, Time: {:.4f}s".format(epoch, avg_train_loss, avg_train_acc*100, avg_valid_loss, avg_valid_acc*100, epoch_end-epoch_start))# Save if the model has best accuracy till now
torch.save(model, dataset+'_model_'+str(epoch)+'.pt')return model, history, best_epoch
# Print the model to be trained
summary(resnet50, input_size=(3,224,224), batch_size=batch_size, device='cuda')# Train the model for 25 epochs# num_epochs = 30# trained_model, history, best_epoch = train_and_validate(resnet50, loss_func, optimizer, num_epochs)# torch.save(history, dataset+'_history.pt')
defcomputeTestSetAccuracy(model, loss_criterion):'''
Function to compute the accuracy on the test set
Parameters
:param model: Model to test
:param loss_criterion: Loss Criterion to minimize
'''
test_acc =0.0
test_loss =0.0# Validation - No gradient tracking needed.with torch.no_grad():# Set to evaluation mode.
model.eval()# Validation loopfor j,(inputs, labels)inenumerate(test_data_loader):
inputs = inputs.to(device)
labels = labels.to(device)# Forward pass - compute outputs on input data using the model.
outputs = model(inputs)# Compute loss
loss = loss_criterion(outputs, labels)# Compute the total loss for the batch and add it to valid_loss.
test_loss += loss.item()* inputs.size(0)# Calculate validation accuracy.
ret, predictions = torch.max(outputs.data,1)
correct_counts = predictions.eq(labels.data.view_as(predictions))# Convert correct_counts to float and then compute the mean.
acc = torch.mean(correct_counts.type(torch.FloatTensor))# Compute total accuracy in the whole batch and add to valid_acc.
test_acc += acc.item()* inputs.size(0)print("Test Batch number: {:03d}, Test: Loss: {:.4f}, Accuracy: {:.4f}".format(j, loss.item(), acc.item()))# Find average test loss and test accuracy
avg_test_loss = test_loss/len(data['test'])
avg_test_acc = test_acc/len(data['test'])print("Test accuracy : "+str(avg_test_acc))
# Get a mapping of the indices to the class names, in order to see the output classes of the test images.
idx_to_class ={v: k for k, v in data['train'].class_to_idx.items()}print(idx_to_class)
defpredict(model, test_image_name):'''
Function to predict the class of a single test image
Parameters
:param model: Model to test
:param test_image_name: Test image
'''
transform = image_transforms['test']
test_image = Image.open(test_image_name)
plt.imshow(test_image)
test_image_tensor = transform(test_image)if torch.cuda.is_available():
test_image_tensor = test_image_tensor.view(1,3,224,224).cuda()else:
test_image_tensor = test_image_tensor.view(1,3,224,224)with torch.no_grad():
model.eval()# Model outputs log probabilities.
out = model(test_image_tensor)
ps = torch.exp(out)
topk, topclass = ps.topk(3, dim=1)
cls = idx_to_class[topclass.cpu().numpy()[0][0]]
score = topk.cpu().numpy()[0][0]for i inrange(3):print("Predcition", i+1,":", idx_to_class[topclass.cpu().numpy()[0][i]],", Score: ", topk.cpu().numpy()[0][i])
# Test a particular model on a test image# ! wget https://cdn.pixabay.com/photo/2018/10/01/12/28/skunk-3716043_1280.jpg -O skunk.jpg
model = torch.load("./Datasets/_model_{}.pt".format(best_epoch))
predict(model,'./Datasets/skunk-3716043_1280.jpg')# Load Data from folders#computeTestSetAccuracy(model, loss_func)