chijioke_ai_proj.py

# import packages
%matplotlib inline
%config InlineBackend.figure_format = 'retina'

from os import path, makedirs
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch import nn, optim
import torch.nn.functional as F
from collections import OrderedDict
from torchvision import datasets, transforms, models
from workspace_utils import active_session


# LOAD THE DATA
data_dir = 'flowers'
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
test_dir = data_dir + '/test'

# Define transforms for the training, validation, and testing sets
normalize_tranform = transforms.Normalize(
    (0.485, 0.456, 0.406),
    (0.229, 0.224, 0.225)
)
rand_flip_v_transform = transforms.RandomVerticalFlip()
rand_flip_h_transform = transforms.RandomHorizontalFlip() # Mirror Transformation
center_crop_transform = transforms.CenterCrop(224) # Cropping from Center
resize_transform = transforms.Resize(225)
rotation_transform = transforms.RandomRotation(180) # from -180 to +180

# Define transforms
data_transforms = {
    'train_transforms': transforms.Compose([
        resize_transform,
        center_crop_transform,
        rand_flip_h_transform,
        rand_flip_v_transform,
        rotation_transform,
        transforms.ToTensor(),
        normalize_tranform
    ]),
    'test_val_transforms': transforms.Compose([
        resize_transform,
        center_crop_transform,
        transforms.ToTensor(),
        normalize_tranform
    ])
}

# Load the datasets with ImageFolder
image_datasets = {
    'train_data': datasets.ImageFolder(train_dir, transform=data_transforms['train_transforms']),
    'test_data': datasets.ImageFolder(test_dir, transform=data_transforms['test_val_transforms']),
    'valid_data': datasets.ImageFolder(valid_dir, transform=data_transforms['test_val_transforms']),
}

# Define the dataloaders, Using the image datasets and the trainforms
data_loaders = {
    'train_loader': torch.utils.data.DataLoader(image_datasets['train_data'], batch_size=17, shuffle=True),
    'test_loader': torch.utils.data.DataLoader(image_datasets['test_data'], batch_size=17, shuffle=True),
    'valid_loader': torch.utils.data.DataLoader(image_datasets['valid_data'], batch_size=17, shuffle=True)
}



# Build and train your network
model = models.vgg13(pretrained=True)
in_features = model.classifier[0].in_features

# Freeze the model's paramerers
for parameter in model.parameters():
    parameter.requires_grad = False
    

# Build out our own classifier
new_classifier = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(in_features, 1000)),
    ('relu1', nn.ReLU()),
    ('dropout1', nn.Dropout(p=0.5)),

    ('hidden_unit', nn.Linear(1000, 102)), # we have 102 labels

    ('output', nn.LogSoftmax(dim=1)),
]))

# Replace the pre-trained network classifier with the new classifier
model.classifier = new_classifier

# Define the device/machine
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# Define Validation method
def validation (model, loader, criterion):
    loss = 0
    accuracy = 0
    
    for images, labels in loader:
        images, labels = images.to(device), labels.to(device)
        output = model.forward(images)
        loss += criterion(output, labels).item()
        
        ps = torch.exp(output) # because we're using NLLLoss() as criterion
#         ps = F.softmax(output, dim=1) # Use this only for nn.CrossEntropyLoss() criterion
        
        equality = (labels.data == ps.max(dim=1)[1])
        accuracy = equality.type(torch.cuda.FloatTensor).mean()
    return loss, accuracy

# Define the Evaluation method
def evaluate_test_data(test_loader):
    total = 0
    correct = 0
    with torch.no_grad():
        for img, lbs in test_loader:
            img, lbs = img.to(device), lbs.to(device)
            out = model.forward(img)
            _, pred = torch.max(out.data, 1)
            total += lbs.size()[0]
            correct += (pred == lbs).sum().item()
    print("Correct: {} ...  \
            Total: {} ...  \
            Accuracy: {}".format(correct, total, (100*correct)/total ))

# Move model to device
model.to(device)
print("Model running on {}".format(device))

# Define optimizer and criterion
optimizer = optim.Adam(model.classifier.parameters(), lr=1e-6)
# criterion = nn.CrossEntropyLoss()
criterion = nn.NLLLoss()

# Define epochs, check_at
epochs = 10
steps = 0
check_at = 30

# Train the network, print out both test and validation loss and accuracy
with active_session():
    for e in range(epochs):
        running_loss = 0
        for images, labels in iter(data_loaders['train_loader']):
            steps += 1
            images, labels = images.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model.forward(images)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()
            
            running_loss += loss.item()
            
            if steps % check_at == 0:
                model.eval()
                print("Epoch: {}/{}...  ".format(e+1, epochs),
                      "Training Loss: {:.4f}".format(running_loss/check_at))
                
                with torch.no_grad():
                    valid_loss, valid_accuracy = validation(model, loader=data_loaders['valid_loader'], criterion=criterion)
                    t_loss, t_accuracy = validation(model, loader=data_loaders['test_loader'], criterion=criterion)
                    
                print("Validation Loss: {:.4f}".format(valid_loss/len(data_loaders['valid_loader'])))
                print("Validation Accuracy: {:.4f}".format(valid_accuracy/len(data_loaders['valid_loader'])))
                # Test data
                print("Test Loss: {:.4f}".format(t_loss/len(data_loaders['test_loader'])))
                print("Test Accuracy: {:.4f}\n".format(t_accuracy/len(data_loaders['test_loader'])))
                
                running_loss = 0
                
                model.train()
    print('Training finished.')

# Evaluate on Test/Training/Validation data
evaluate_test_data(data_loaders['train_loader'])
evaluate_test_data(data_loaders['test_loader'])
evaluate_test_data(data_loaders['valid_loader'])

# # Save Checkpoint
# checkpoint = {
#     'architecture': 'vgg13',
#     'n_in_features': model.classifier[0].in_features,
#     'n_out_features': model.classifier[0].out_features,
#     'hidden_units': 512,
#     'model_state_dict': model.state_dict(),
# #     'optimizer_state_dict': optimizer.state_dict(),
#     'epochs': epochs,
#     'class_to_idx': image_datasets['test_data'].class_to_idx
# }

# checkpoint_path = 'checkpoints_dir'
# if not path.exists(checkpoint_path):
#     makedirs('./' + checkpoint_path)
# checkpoint_path = path.join('checkpoints_dir', '__checkpoint.pth')
# torch.save(checkpoint, checkpoint_path)

# # Define the method to load the checkpoint from file_path
# def load_checkpoint(filepath):
#     return torch.load(filepath, map_location=lambda storage, loc: storage)

# # Load the checkpoint
# checkpoint = load_checkpoint('./checkpoints_dir/__checkpoint.pth')

# loaded_classifier = nn.Sequential(OrderedDict([

#     ('fc1', nn.Linear(checkpoint['n_in_features'], 1000)),
#     ('relu1', nn.ReLU()),
#     ('dropout1', nn.Dropout(p=0.5)),

#     ('hidden_unit', nn.Linear(1000, 102)), # use the values in the checkpoint

#     ('output', nn.LogSoftmax(dim=1)),
# ]))

# # reload model and replace classifier
# model.classifier = loaded_classifier
# model.load_state_dict(checkpoint['model_state_dict'])