PyTorchによる画像分類(vgg16編)

# Imports / Requirements
import json
import numpy as np
import torch
import matplotlib.pyplot as plt
import torch.nn.functional as F
import torchvision
from torch import nn, optim
from torchvision import datasets, transforms, models
from torch.autograd import Variable
from collections import OrderedDict
from PIL import Image

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

!python --version
print(f"PyTorch Version {torch.__version__}")
%matplotlib inline
plt.rcParams.update({'font.size': 22, 'font.family': 'STIXGeneral', 'mathtext.fontset': 'stix'})

%download https://s3.amazonaws.com/content.udacity-data.com/nd089/flower_data.tar.gz

!mkdir flowers

cd flowers

! mv ../flower_data.tar.gz .

!tar -xvzf flower_data.tar.gz

ls

!rm flower_data.tar.gz

cd ..

data_dir = 'flowers'
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
test_dir = data_dir + '/test'

# pre-trained network expectations
# see: https://pytorch.org/docs/stable/torchvision/models.html
expected_means = [0.485, 0.456, 0.406]
expected_std = [0.229, 0.224, 0.225]
max_image_size = 224
batch_size = 32

# DONE: Define your transforms for the training, validation, and testing sets
data_transforms = {
    "training": transforms.Compose([transforms.RandomHorizontalFlip(p=0.25),
                                    transforms.RandomRotation(25),
                                    transforms.RandomGrayscale(p=0.02),
                                    transforms.RandomResizedCrop(max_image_size),
                                    transforms.ToTensor(),
                                    transforms.Normalize(expected_means, expected_std)]),

    "validation": transforms.Compose([transforms.Resize(max_image_size + 1),
                                      transforms.CenterCrop(max_image_size),
                                      transforms.ToTensor(),
                                      transforms.Normalize(expected_means, expected_std)]),

    "testing": transforms.Compose([transforms.Resize(max_image_size + 1),
                                   transforms.CenterCrop(max_image_size),
                                   transforms.ToTensor(),
                                   transforms.Normalize(expected_means, expected_std)])

}

# DONE: Load the datasets with ImageFolder
image_datasets = {
    "training": datasets.ImageFolder(train_dir, transform=data_transforms["training"]),
    "validation": datasets.ImageFolder(valid_dir, transform=data_transforms["validation"]),
    "testing": datasets.ImageFolder(test_dir, transform=data_transforms["testing"])
}

# DONE: Using the image datasets and the trainforms, define the dataloaders
dataloaders = {
    "training": torch.utils.data.DataLoader(image_datasets["training"], batch_size=batch_size, shuffle=True),
    "validation": torch.utils.data.DataLoader(image_datasets["validation"], batch_size=batch_size),
    "testing": torch.utils.data.DataLoader(image_datasets["testing"], batch_size=batch_size)
}

import json

with open('cat_to_name.json', 'r') as f:
    cat_to_name = json.load(f)

print(f"Images are labeled with {len(cat_to_name)} categories.")

! find / -name cat_to_name.json

%download https://raw.githubusercontent.com/cjimti/aipnd-project/master/cat_to_name.json -f flowers/cat_to_name.json

ls flowers

import json

with open('flowers/cat_to_name.json', 'r') as f:
    cat_to_name = json.load(f)

print(f"Images are labeled with {len(cat_to_name)} categories.")

# DONE: Build and train your network
# Get model Output Size = Number of Categories
output_size = len(cat_to_name)

# Using VGG16.
nn_model = models.vgg16(pretrained=True)

# Input size from current classifier
input_size = nn_model.classifier[0].in_features
hidden_size = [
    (input_size // 8),
    (input_size // 32)
]

# Prevent backpropigation on parameters
for param in nn_model.parameters():
    param.requires_grad = False 
    
# Create nn.Module with Sequential using an OrderedDict
# See https://pytorch.org/docs/stable/nn.html#torch.nn.Sequential
classifier = nn.Sequential(OrderedDict([
        ('fc1', nn.Linear(input_size, hidden_size[0])),
        ('relu1', nn.ReLU()),
        ('dropout', nn.Dropout(p=0.15)),
        ('fc2', nn.Linear(hidden_size[0], hidden_size[1])),
        ('relu2', nn.ReLU()),
        ('dropout', nn.Dropout(p=0.15)),
        ('output', nn.Linear(hidden_size[1], output_size)),
        # LogSoftmax is needed by NLLLoss criterion
        ('softmax', nn.LogSoftmax(dim=1))
    ]))
    
# Replace classifier
nn_model.classifier = classifier

hidden_size

torch.cuda.is_available()

device

# hyperparameters
# https://en.wikipedia.org/wiki/Hyperparameter
epochs = 5
learning_rate = 0.001
chk_every = 50

# Start clean by setting gradients of all parameters to zero. 
nn_model.zero_grad()

# The negative log likelihood loss as criterion.
criterion = nn.NLLLoss()

# Adam: A Method for Stochastic Optimization
# https://arxiv.org/abs/1412.6980
optimizer = optim.Adam(nn_model.classifier.parameters(), lr=learning_rate)

# Move model to perferred device.
nn_model = nn_model.to(device)

data_set_len = len(dataloaders["training"].batch_sampler)
total_val_images = len(dataloaders["validation"].batch_sampler) * dataloaders["validation"].batch_size

print(f'Using the {device} device to train.')
print(f'Training on {data_set_len} batches of {dataloaders["training"].batch_size}.')
print(f'Displaying average loss and accuracy for epoch every {chk_every} batches.')

for e in range(epochs):
    e_loss = 0
    prev_chk = 0
    total = 0
    correct = 0
    print(f'\nEpoch {e+1} of {epochs}\n----------------------------')
    for ii, (images, labels) in enumerate(dataloaders["training"]):
        # Move images and labeles preferred device
        # if they are not already there
        images = images.to(device)
        labels = labels.to(device)
        
        # Set gradients of all parameters to zero. 
        optimizer.zero_grad()
        
        # Propigate forward and backward 
        outputs = nn_model.forward(images)
        
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        
        # Keep a running total of loss for
        # this epoch
        e_loss += loss.item()
        
        # Accuracy
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        # Keep a running total of loss for
        # this epoch
        itr = (ii + 1)
        if itr % chk_every == 0:
            avg_loss = f'avg. loss: {e_loss/itr:.4f}'
            acc = f'accuracy: {(correct/total) * 100:.2f}%'
            print(f'  Batches {prev_chk:03} to {itr:03}: {avg_loss}, {acc}.')
            prev_chk = (ii + 1)
    
    # Validate Epoch
    e_valid_correct = 0
    e_valid_total = 0    

    # Disabling gradient calculation
    with torch.no_grad():
        for ii, (images, labels) in enumerate(dataloaders["validation"]):
            # Move images and labeles perferred device
            # if they are not already there
            images = images.to(device)
            labels = labels.to(device)

            outputs = nn_model(images)
            _, predicted = torch.max(outputs.data, 1)
            e_valid_total += labels.size(0)
            e_valid_correct += (predicted == labels).sum().item()
        print(f"\n\tValidating for epoch {e+1}...")
        correct_perc = 0
        if e_valid_correct > 0:
            correct_perc = (100 * e_valid_correct // e_valid_total)
        print(f'\tAccurately classified {correct_perc:d}% of {total_val_images} images.')

print('Done...')

# DONE: Do validation on the test set
correct = 0
total = 0
total_images = len(dataloaders["testing"].batch_sampler) * dataloaders["testing"].batch_size

# Disabling gradient calculation
with torch.no_grad():
    for ii, (images, labels) in enumerate(dataloaders["testing"]):
        # Move images and labeles perferred device
        # if they are not already there
        images = images.to(device)
        labels = labels.to(device)
        
        outputs = nn_model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

print(f'Accurately classified {(100 * correct // total):d}% of {total_images} images.')

# DONE: Save the checkpoint
def save_checkpoint(model_state, file='checkpoint.pth'):
    torch.save(model_state, file)

nn_model.class_to_idx = image_datasets['training'].class_to_idx
model_state = {
    'epoch': epochs,
    'state_dict': nn_model.state_dict(),
    'optimizer_dict': optimizer.state_dict(),
    'classifier': classifier,
    'class_to_idx': nn_model.class_to_idx,
}

save_checkpoint(model_state, 'checkpoint.pth')

# DONE: Write a function that loads a checkpoint and rebuilds the model

def load_checkpoint(file='checkpoint.pth'):
    # Loading weights for CPU model while trained on GP
    # https://discuss.pytorch.org/t/loading-weights-for-cpu-model-while-trained-on-gpu/1032
    model_state = torch.load(file, map_location=lambda storage, loc: storage)
    model = models.vgg16(pretrained=True)
    model.classifier = model_state['classifier']
    model.load_state_dict(model_state['state_dict'])
    model.class_to_idx = model_state['class_to_idx']
    return model

chkp_model = load_checkpoint()

!ls -lh checkpoint.pth

def process_image(image):
    ''' Scales, crops, and normalizes a PIL image for a PyTorch model,
        returns an Numpy array
    '''
    expects_means = [0.485, 0.456, 0.406]
    expects_std = [0.229, 0.224, 0.225]           
    pil_image = Image.open(image).convert("RGB")
    
    # Any reason not to let transforms do all the work here?
    in_transforms = transforms.Compose([transforms.Resize(256),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize(expects_means, expects_std)])
    pil_image = in_transforms(pil_image)
    return pil_image

# DONE: Process a PIL image for use in a PyTorch model
chk_image = process_image(valid_dir + '/1/image_06739.jpg')
type(chk_image)

def imshow(image, ax=None, title=None):
    if ax is None:
        fig, ax = plt.subplots()
    
    # PyTorch tensors assume the color channel is the first dimension
    # but matplotlib assumes is the third dimension
    image = image.transpose((1, 2, 0))
    
    # Undo preprocessing
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    image = std * image + mean
    
    # Image needs to be clipped between 0 and 1 or it looks like noise when displayed
    image = np.clip(image, 0, 1)    
    ax.imshow(image)    
    return ax

imshow(chk_image.numpy())

def predict(image_path, model, topk=5):
    ''' Predict the class (or classes) of an image using a trained deep learning model.
    '''
    # DONE: Implement the code to predict the class from an image file
    
    # evaluation mode
    # https://pytorch.org/docs/stable/nn.html#torch.nn.Module.eval
    model.eval()
    
    # cpu mode
    model.cpu()
    
    # load image as torch.Tensor
    image = process_image(image_path)
    
    # Unsqueeze returns a new tensor with a dimension of size one
    # https://pytorch.org/docs/stable/torch.html#torch.unsqueeze
    image = image.unsqueeze(0)
    
    # Disabling gradient calculation 
    # (not needed with evaluation mode?)
    with torch.no_grad():
        output = model.forward(image)
        top_prob, top_labels = torch.topk(output, topk)
        
        # Calculate the exponentials
        top_prob = top_prob.exp()
        
    class_to_idx_inv = {model.class_to_idx[k]: k for k in model.class_to_idx}
    mapped_classes = list()
    
    for label in top_labels.numpy()[0]:
        mapped_classes.append(class_to_idx_inv[label])
        
    return top_prob.numpy()[0], mapped_classes

# DONE: Display an image along with the top 5 classes
chk_image_file = valid_dir + '/55/image_04696.jpg'
correct_class = cat_to_name['55']

top_prob, top_classes = predict(chk_image_file, chkp_model)

label = top_classes[0]

fig = plt.figure(figsize=(16,16))
sp_img = plt.subplot2grid((15,9), (0,0), colspan=9, rowspan=9)
sp_prd = plt.subplot2grid((15,9), (9,2), colspan=5, rowspan=5)

image = Image.open(chk_image_file)
sp_img.axis('off')
sp_img.set_title(f'{cat_to_name[label]}')
sp_img.imshow(image)

labels = []
for class_idx in top_classes:
    labels.append(cat_to_name[class_idx])

yp = np.arange(5)
sp_prd.set_yticks(yp)
sp_prd.set_yticklabels(labels)
sp_prd.set_xlabel('Probability')
sp_prd.invert_yaxis()
sp_prd.barh(yp, top_prob, xerr=0, align='center', color='blue')

plt.show()
print(f'Correct classification: {correct_class}')
print(f'Correct prediction: {correct_class == cat_to_name[label]}')