Spaces:
Build error
Build error
| """ | |
| In this TorchDaily we will TRAIN | |
| A MODEL USING TRANSFER LEARNING | |
| Cats Vs Dogs Dataset | |
| EARLIER ACC==14% OR LESS | |
| NOW ITS 70% AND MORE | |
| THE POWER OF ALEXNET (PRETRAINED MODELS IS VISIBLE) | |
| DATE ==> 10-05-21 | |
| """ | |
| import torch | |
| import torch.nn as nn | |
| from torch.utils.data import DataLoader | |
| import matplotlib.pyplot as plt | |
| from torchvision import transforms, datasets, models | |
| import torchvision | |
| from tqdm import tqdm | |
| import os | |
| import PIL.Image as Image | |
| import time | |
| import torch, torchvision | |
| from torchvision import datasets, models, transforms | |
| import torch.nn as nn | |
| import torch.optim as optim | |
| from torch.utils.data import DataLoader | |
| import time | |
| # from torchsummary import summary | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import os | |
| from PIL import Image | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| print(device) | |
| # prepare data | |
| convert = transforms.Compose( | |
| [ | |
| transforms.Resize((128, 128)), | |
| transforms.RandomHorizontalFlip(0.2), | |
| transforms.ToTensor(), | |
| ] | |
| ) | |
| # dataloader | |
| data = datasets.ImageFolder(root="PetImages/", transform=convert) | |
| Loader = DataLoader(data, batch_size=64, shuffle=True) | |
| MAP = {0: "Cat", 1: "Dog"} | |
| ##UNCOMMENT FOR SEEING THE DATA IMAGES | |
| # fig, ax = plt.subplots(8, 8, figsize=(20, 20)) | |
| # fig.suptitle("Dogs And Cats IMages") | |
| # for i, (img, lab) in zip(range(0, 8 * 8), Loader): | |
| # x = i // 8 | |
| # y = i % 8 | |
| # print(f"{x},{y}") | |
| # ax[x, y].imshow(img[i].squeeze().permute(1,2,0)) | |
| # ax[x, y].set_title(f"{lab[i]}") | |
| # ax[x, y].axis("off") | |
| # plt.show() | |
| # # Add on classifier | |
| # # HOW TO CHANGE THE INPUT LAYER WHICH ACCEPTS THE 224*224 INPUT | |
| # # I WANNA CHANGE THAT TO 128*128 THIS SIZE WILL SUFFICE | |
| # We Use AlexNet for transfer learning | |
| ##answers below | |
| alexnet = torchvision.models.alexnet(pretrained=True) | |
| for param in alexnet.parameters(): | |
| param.requires_grad = False | |
| # Add a avgpool here | |
| avgpool = nn.AdaptiveAvgPool2d((7, 7)) | |
| # Replace the classifier layer | |
| # to customise it according to our output | |
| alexnet.classifier = nn.Sequential( | |
| nn.Linear(256 * 7 * 7, 1024), | |
| nn.Linear(1024, 256), | |
| nn.Linear(256, 2), | |
| ) | |
| # putting model in a training mode | |
| alexnet.train() | |
| print(alexnet) | |
| criterion = nn.CrossEntropyLoss() | |
| optimizer = torch.optim.Adam(alexnet.parameters(), lr=0.001) | |
| EPOCHS = 4 | |
| TRAIN = False | |
| losses = [] | |
| def train_and_validate(model, loss_criterion, optimizer, epochs=25): | |
| """ | |
| Function 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) | |
| Returns | |
| model: Trained Model with best validation accuracy | |
| history: (dict object): Having training loss, accuracy and validation loss, accuracy | |
| """ | |
| start = time.time() | |
| history = [] | |
| best_acc = 0.0 | |
| for epoch in range(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.0 | |
| for i, (inputs, labels) in enumerate(Loader): | |
| inputs = inputs.to(device) | |
| labels = labels.to(device) | |
| # Clean existing gradients | |
| optimizer.zero_grad() | |
| # Forward pass - compute outputs on input data using the model | |
| x = alexnet.features(inputs) | |
| x = avgpool(x) | |
| x = x.view(-1, 256 * 7 * 7) | |
| outputs = alexnet.classifier(x) | |
| # Compute loss | |
| loss = loss_criterion(outputs, labels) | |
| # Backpropagate the gradients | |
| loss.backward() | |
| # Update the parameters | |
| 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 | |
| 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 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 j, (inputs, labels) in enumerate(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 | |
| 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 | |
| valid_acc += acc.item() * inputs.size(0) | |
| # print("Validation Batch number: {:03d}, Validation: Loss: {:.4f}, Accuracy: {:.4f}".format(j, loss.item(), acc.item())) | |
| # Find average training loss and training accuracy | |
| avg_train_loss = train_loss / train_data_size | |
| avg_train_acc = train_acc / train_data_size | |
| # Find average training loss and training accuracy | |
| avg_valid_loss = valid_loss / valid_data_size | |
| avg_valid_acc = valid_acc / valid_data_size | |
| 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 + 1, | |
| 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, "CatVsDogsModel.pth") | |
| return model, history | |
| if TRAIN: | |
| trained_model, history = train_and_validate(alexnet, criterion, optimizer, EPOCHS) | |
| plt.plot(losses) | |
| plt.show() | |
| history = np.array(history) | |
| plt.plot(history[:, 0:2]) | |
| plt.legend(["Tr Loss", "Val Loss"]) | |
| plt.xlabel("Epoch Number") | |
| plt.ylabel("Loss") | |
| plt.ylim(0, 1) | |
| plt.savefig(dataset + "_loss_curve.png") | |
| plt.show() | |
| plt.plot(history[:, 2:4]) | |
| plt.legend(["Tr Accuracy", "Val Accuracy"]) | |
| plt.xlabel("Epoch Number") | |
| plt.ylabel("Accuracy") | |
| plt.ylim(0, 1) | |
| plt.savefig(dataset + "_accuracy_curve.png") | |
| plt.show() | |
| TEST = False | |
| history = [] | |
| def test(): | |
| test = datasets.ImageFolder(root="PetTest/", transform=convert) | |
| testLoader = DataLoader(test, batch_size=16, shuffle=False) | |
| checkpoint = torch.load("CatVsDogsModel.pth") | |
| alexnet.load_state_dict(checkpoint["state_dict"]) | |
| optimizer.load_state_dict(checkpoint["optimizer"]) | |
| for params in alexnet.parameters(): | |
| params.requires_grad == False | |
| print(alexnet) | |
| with torch.no_grad(): | |
| # Set to evaluation mode | |
| alexnet.eval() | |
| train_data_size = 101 | |
| valid_data_size = 101 | |
| # Validation loop | |
| # Loss and Accuracy within the epoch | |
| valid_loss = 0.0 | |
| valid_acc = 0.0 | |
| for j, (inputs, labels) in enumerate(testLoader): | |
| inputs = inputs.to(device) | |
| labels = labels.to(device) | |
| # Forward pass - compute outputs on input data using the model | |
| x = alexnet.features(inputs) | |
| x = avgpool(x) | |
| x = x.view(-1, 256 * 7 * 7) | |
| outputs = alexnet.classifier(x) | |
| # Compute 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 | |
| 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 | |
| valid_acc += acc.item() * inputs.size(0) | |
| print( | |
| """Validation Batch number: {:03d}, | |
| Validation: Loss: {:.4f}, | |
| Accuracy: {:.4f}""".format( | |
| j, loss.item(), acc.item() | |
| ) | |
| ) | |
| # Find average training loss and training accuracy | |
| avg_valid_loss = valid_loss / valid_data_size | |
| avg_valid_acc = valid_acc / valid_data_size | |
| history.append([avg_valid_loss, avg_valid_acc]) | |
| print( | |
| " Training: Loss: {:.4f}, Accuracy: {:.4f}%, \n\t\tValidation : Loss : {:.4f}, Accuracy: {:.4f}%".format( | |
| avg_train_loss, | |
| avg_train_acc * 100, | |
| avg_valid_loss, | |
| avg_valid_acc * 100, | |
| ) | |
| ) | |
| plt.plot(valid_acc) | |
| plt.plot(valid_loss) | |
| plt.show() | |
| if TEST: | |
| test() | |
| print("Validation Complete") | |
| with open("ModelHistory.txt", "w") as f: | |
| for i in history: | |
| f.writelines(f"{i}") | |
| print("Validation Complete") | |
| ## This model reported a accuracy of 97%(on DOGS ONLY) using AlexNet | |
| ## the Pros of using a pretrained model is clearly seen here | |
| ## date -- 13th April 2021 (thursday) | |
| ####ACCURACY AND OTHER THINGS TOO TO BE APPENDED SOON ###### | |
| PREDICT = True | |
| def predict(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 | |
| """ | |
| # try: | |
| transform = transforms.Compose( | |
| [transforms.Resize((128, 128)), transforms.ToTensor()] | |
| ) | |
| test_image = Image.open(test_image_name) | |
| test_image_tensor = transform(test_image).to(device) | |
| plt.imshow(test_image) | |
| plt.axis("off") | |
| plt.imshow(test_image_tensor.cpu().squeeze().permute(1, 2, 0)) | |
| plt.show() | |
| with torch.no_grad(): | |
| model.eval() | |
| test_image_tensor = test_image_tensor.unsqueeze(0) | |
| print(test_image_tensor.shape) | |
| x = alexnet.features(test_image_tensor) | |
| x = avgpool(x) | |
| x = x.view(-1, 256 * 7 * 7) | |
| out = alexnet.classifier(x) | |
| ###THESE ARE SCORES OF THE ACC. ### | |
| ### UNCOMMENT TO SEE THE SCORES OF EACH CLASS ### | |
| # ps = torch.exp(out) | |
| # print(f'ps: {ps}') | |
| # topk, topclass = ps.topk(2, dim=1) | |
| # print(f'ps.topk: {ps.topk(2, dim=1)}') | |
| # print(f'topclass: {topclass}') | |
| print("Predcition", MAP[out.numpy().argmax()]) | |
| # print(f"out: {out.numpy().argmax()}") | |
| # except Exception as e: | |
| # print(e) | |
| if PREDICT: | |
| checkpoint = torch.load( | |
| "CatVsDogsModel.pth", map_location=torch.device("cpu") | |
| ) | |
| alexnet.load_state_dict(checkpoint["state_dict"]) | |
| alexnet = alexnet.to(device) | |
| optimizer.load_state_dict(checkpoint["optimizer"]) | |
| for params in alexnet.parameters(): | |
| params.requires_grad == False | |
| print(predict(alexnet, "PetTest/Cat/12401.jpg")) | |