AI challenger 场景分类 PyTorch 迁移学习 resnet18

没计算资源，只能简单测试下resnet18

训练代码（带验证）：

'''
TODO:
- 采用更深的网络（简单，但是需要计算资源）
- top3 accuracy，可参考https://github.com/pytorch/examples/blob/master/imagenet/main.py
'''
#pkill -9 python
#nvidia-smi
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.autograd import Variable
import pandas as pd
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import transforms, utils
import time
import json
#plt.ion()   # interactive mode
%matplotlib inline

with open('../ai_challenger_scene_train_20170904/scene_train_annotations_20170904.json', 'r') as f: #label文件
    label_raw_train = json.load(f)
with open('../ai_challenger_scene_validation_20170908/scene_validation_annotations_20170908.json', 'r') as f: #label文件
    label_raw_val = json.load(f)

label_raw_train[0]['label_id']
len(label_raw_train)

class SceneDataset(Dataset):

    def __init__(self, json_labels, root_dir, transform=None):
        """
        Args:
            json_labesl (list):read from official json file.
            root_dir (string): Directory with all the images.
            transform (callable, optional): Optional transform to be applied
                on a sample.
        """
        self.label_raw = json_labels
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return len(self.label_raw)

    def __getitem__(self, idx):
        img_name = os.path.join(self.root_dir, self.label_raw[idx]['image_id'])
        image = Image.open(img_name)
        label = int(self.label_raw[idx]['label_id'])

        if self.transform:
            image = self.transform(image)

        return image, label

data_transforms = {
    'train': transforms.Compose([
        transforms.RandomSizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Scale(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}
transformed_dataset_train = SceneDataset(json_labels=label_raw_train,
                                    root_dir='../ai_challenger_scene_train_20170904/scene_train_images_20170904',
                                           transform=data_transforms['train']
                                           )      
transformed_dataset_val = SceneDataset(json_labels=label_raw_val,
                                    root_dir='../ai_challenger_scene_validation_20170908/scene_validation_images_20170908',
                                           transform=data_transforms['val']
                                           )         
batch_size = 64
dataloader = {
  'train':DataLoader(transformed_dataset_train, batch_size=batch_size,shuffle=True, num_workers=8),
             'val':DataLoader(transformed_dataset_val, batch_size=batch_size,shuffle=True, num_workers=8)
             }
dataset_sizes = {
  'train': len(label_raw_train), 'val':len(label_raw_val)}
use_gpu = torch.cuda.is_available()
#use_gpu = False

def imshow(inp, title=None):
    """Imshow for Tensor."""
    inp = inp.numpy().transpose((1, 2, 0))
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    inp = std * inp + mean
    plt.imshow(inp)
    if title is not None:
        plt.title(title)
    plt.pause(0.001)  # pause a bit so that plots are updated


# Get a batch of training data
inputs, classes = next(iter(dataloader['train']))

# Make a grid from batch
out = torchvision.utils.make_grid(inputs)

imshow(out)


######################################################################
# Training the model
# ------------------
#
# Now, let's write a general function to train a model. Here, we will
# illustrate:
#
# -  Scheduling the learning rate
# -  Saving the best model
#
# In the following, parameter ``scheduler`` is an LR scheduler object from
# ``torch.optim.lr_scheduler``.


def train_model (model, criterion, optimizer, scheduler, num_epochs, total_steps):
    since = time.time()

    print('total_steps is %d' % total_steps)
    mystep = 0

    best_model_wts = model.state_dict()
    best_acc = 0.0

    for epoch in range(num_epochs):
        print('Epoch {}/{}'.format(epoch, num_epochs - 1))
        print('-' * 10)


        if (epoch%10 == 0):
            torch.save(best_model_wts, ('resnet18_model_wts_%d.pth')% epoch)

        # Each epoch has a training and validation phase
        for phase in ['train', 'val'