【学习笔记】PyTorch（2）—Kaggle竞赛树叶分类

导入包

import pandas as pd
import numpy as np
import collections
import matplotlib.pyplot as plt
from PIL import Image

import torch
from torch import nn
from torchvision import transforms, models
from torch.utils.data import Dataset, DataLoader
import torchvision.models as models

探索数据

train_path = 'train.csv'
train_data = pd.read_csv(train_path, header=0)
print(train_data.head())
train_data.describe()
num_unique = train_data['label'].unique()
num_label = len(num_unique)# 分为176类
num_label

制造数据集

train_label = sorted(list(set(train_data['label'])))
num_label = len(train_label)
#名称转化为数值标签
label_to_num = dict(zip(train_label, range(num_label)))
#数值标签转化为名称
num_to_label = dict(zip( range(num_label),train_label))
#测试一下
print(num_to_label[0])
print(label_to_num[num_to_label[0]])

产生LeavesDataset数据集

class LeavesDataset():
    def __init__(self, csv_path, img_path, mode='train'):
        """
        Args:
            csv_path (string): csv 文件路径
            img_path (string): 图像文件所在路径
            mode (string): 训练模式还是测试模式
        """
        self.mode = mode
        data_info = pd.read_csv(csv_path)
        data_len = len(data_info)
        self.img_path = img_path

        if mode == 'test':
            image_arr = np.array(data_info.iloc[:,0])
            self.image_arr = image_arr
        else:
            # 验证集：9 / 10；训练集：1 / 10
            if mode == 'train':
                indices = [i for i in range(data_len) if i % 10 != 0]
            elif mode == 'valid':
                indices = [i for i in range(data_len) if i % 10 == 0]
            image_arr = np.array(data_info.iloc[indices,0])
            label_arr = np.array(data_info.iloc[indices,1])
            self.image_arr = image_arr
            self.label_arr = label_arr
        self.real_len = len(self.image_arr)
        print('Finished reading the {} set of the LeavesDataset ({} samples found)'.format(mode, self.real_len))
        
    # 重写__getitem__函数
    def __getitem__(self,index):
        assert index <= len(self)
        single_image_name = self.image_arr[index]
        image = Image.open(self.img_path + single_image_name)
        if self.mode == 'train':
            transform=transforms.Compose(
                [transforms.Resize((224,224)),
                transforms.RandomHorizontalFlip(p=0.5),#随机水平翻转
                transforms.RandomVerticalFlip(p=0.5),#随机垂直翻转
                transforms.ToTensor()
                ])
        else:
            # valid test 不做数据增强
            transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
        image = transform(image)

        if self.mode == 'test':
            return image
        else:
            label = self.label_arr[index]
            number_label = label_to_num[label]
            return image,number_label
    def __len__(self):
        return self.real_len

加载数据

def prep_dataloader(train_path, img_path, mode, batch_size):
    dataset = LeavesDataset(train_path, img_path, mode)
    dataloader = DataLoader(dataset, batch_size, shuffle=(mode == 'train'), drop_last=False)
    return dataloader

定义模型

• 冻住预训练模型的一些层

# 不需要训练的层
def set_parameter_requires_grad(model, feature_extracting):
    '''
    冻住层，将对应的参数梯度设为零，不能反向传播
    '''
    if feature_extracting:
        for param in model.parameters():
            param.requires_grad =False

• gpu是否可用

def get_device():
    return 'cuda' if torch.cuda.is_available() else 'cpu'

• 模型初始化
  使用resnet18网络架构

def initialize_model(num_label, feature_extract = False, use_pretrained=True):
    model_ft = models.resnet18(pretrained = use_pretrained)
    set_parameter_requires_grad(model_ft,feature_extract)
    num_ft = model_ft.fc.in_features
    # num_ft = model_ft.fc.in_features
    model_ft.fc = nn.Sequential(nn.Linear(num_ft, num_label))
    return model_ft

• 画学习曲线函数

def plot_learning_curve(loss_record):
    total_steps = len(loss_record['train'])
    x = range(1,total_steps + 1)
    plt.figure(figsize=(6,4))
    plt.plot(x, loss_record['train'], 'r' ,label='train')
    plt.plot(x, loss_record['valid'], 'g', label='valid')
    plt.xlabel('Training epoch')
    plt.ylabel('Loss')
    plt.title('Learning curve')
    plt.legend()
    plt.show()

train

def train(train_path, img_path, model, batch_size, epoch):
    # 数据集
    train_dataset = prep_dataloader(train_path, img_path, mode = 'train', batch_size=batch_size)
    valid_dataset = prep_dataloader(train_path, img_path, mode = 'valid', batch_size=batch_size)
    device = get_device()
    # 损失函数
    criterion = nn.CrossEntropyLoss()
    loss_record = {'train':[], 'valid':[]}
    # 优化器
    learning_rate = 0.0001
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
    criterion = criterion.to(device)
    model = model.to(device)

    # 训练开始
    for i in range(epoch):
        print("--------------第 {} 轮训练开始-----------------".format(i+1))
        train_step = 0
        valid_step = 0
        a_epoch_train_loss = 0
        a_epoch_valid_loss = 0
        #训练
        model.train()
        for img, label in train_dataset:
            img, label = img.to(device), label.to(device)
            train_label = model(img)
            train_loss = criterion(train_label, label)
            optimizer.zero_grad()
            train_loss.backward()
            optimizer.step()
            train_step += 1
            a_epoch_train_loss += train_loss
            a_epoch_train_loss_averge = a_epoch_train_loss / train_step
        loss_record['train'].append(a_epoch_train_loss_averge.detach().cpu().item())

    # 测试开始
        model.eval()
        with torch.no_grad():
            for img, label in valid_dataset:
                img, label = img.to(device), label.to(device)
                train_label = model(img)
                valid_loss = criterion(train_label, label)
                valid_step += 1
                a_epoch_valid_loss += valid_loss
                a_epoch_train_loss_averge = a_epoch_valid_loss / valid_step
            loss_record['valid'].append(a_epoch_train_loss_averge.detach().cpu().item())
    torch.save(model.state_dict(),"model.pth")
    return loss_record

开始训练

• 设置参数开始训练

train_path = '../classify-leaves/train.csv'
img_path = '../classify-leaves/'
model = initialize_model(num_label, feature_extract = False)
# model = Model()
# model = model.to(get_device())
batch_size = 32
epoch = 20
loss_record = train(train_path, img_path, model, batch_size, epoch)

• 画学习曲线

# 画学习曲线
plot_learning_curve(loss_record)

简单测试

model = initialize_model(num_label, feature_extract = False)
model = model.to(get_device())
model.load_state_dict(torch.load('model.pth'))
valid_dataset = prep_dataloader(train_path, img_path, mode = 'valid', batch_size=batch_size)
model.eval()
#  生成迭代器
dataiter = iter(valid_dataset)
dataiter.next()
img, label = dataiter.next()
img, label = img.to(get_device()), label.to(get_device())
train_label = model(img).argmax(dim=1).to('cpu').numpy()
label.to('cpu').numpy() == train_label.T

预测

# 预测
model.eval()
pred = []
test_path = 'test.csv'
test_dataset = prep_dataloader(test_path, img_path, mode = 'test', batch_size=batch_size)
for img in test_dataset:
    img = img.to(get_device())
    with torch.no_grad():
        test_label = model(img)
    pred.extend(test_label.argmax(-1).cpu().numpy().tolist())
pred_to_str = [num_to_label[i] for i in pred]
test_data = pd.read_csv(test_path)
test_data['label'] = pred_to_str
test_data.to_csv('submission.csv', index=False)