pytorch迁移学习+ResNet50实现猫十二分类——高准确率版

之前写过一篇实现猫十二分类的文章，写出了大体的流程，但实际效果并不佳。本文采取微调整个预训练模型的方式，使准确率从0.3提升到了0.93。大体流程参考ResNet猫十二分类，本文只给出不同的地方。

1. 数据

首先数据部分有一些改动

img_datasets = {x: myData(x) for x in ['train', 'val', 'test']}
dataset_sizes = {x: len(img_datasets[x]) for x in ['train', 'val', 'test']}

# dataset准备完毕，开始创建dataloader
train_loader = DataLoader(
    dataset=img_datasets['train'],
    batch_size=batches,
    shuffle=True
)

val_loader = DataLoader(
    dataset=img_datasets['val'],
    batch_size=1,
    shuffle=False
)

test_loader = DataLoader(
    dataset=img_datasets['test'],
    batch_size=1,
    shuffle=False
)

dataloaders = {
    'train': train_loader,
    'val': val_loader,
    'test': test_loader
}

2. 训练

然后是训练部分。训练部分采用每一个epoch中都训练train和val，根据不同epoch中val的准确率的不同不断更新模型的最佳参数，最后训练结束时返回最佳模型参数

# train
# 模型训练的函数 ##################################################
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def train_model(model, criterion, optimizer, scheduler, num_epochs=25):

    best_model_wts = copy.deepcopy(model.state_dict()) 
    best_acc = 0.0  # 最佳正确率，用于判断是否替换best_model_wts

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

        for phase in ['train', 'val']:  # 每个epoch中都包含训练与验证两个阶段
            if phase == 'train':  
                model.train() 
            else:  # 测试阶段
                model.eval() 

            running_loss = 0.0
            running_corrects = 0
            for inputs, labels in dataloaders[phase]:

                inputs = inputs.to(device)  
                labels = labels.to(device)

                optimizer.zero_grad()  

                with torch.set_grad_enabled( phase == 'train'):  
                    outputs = model(inputs)  
                    _, preds = torch.max(outputs, 1) 
                    loss = criterion(outputs, labels)  

                    if phase == 'train':
                        loss.backward()  
                        optimizer.step() 

                running_loss += loss.item() * inputs.size(0)  
                running_corrects += torch.sum(preds == labels.data)  
            if phase == 'train': 
                scheduler.step()  

            epoch_loss = running_loss / dataset_sizes[phase] 
            epoch_acc = running_corrects.double() / dataset_sizes[phase] 

            print('{} Loss: {:.4f} Acc: {:.4f}'.format(  
                phase, epoch_loss, epoch_acc))

            if phase == 'val' and epoch_acc > best_acc:  
                best_acc = epoch_acc  
                best_model_wts = copy.deepcopy(model.state_dict()) 

        print()  

    print('Best val Acc: {:4f}'.format(best_acc)) 

    model.load_state_dict(best_model_wts)
    return model

3. 迁移学习

使用pretrain = True的方式得到预训练模型，更改全连接层的输出维度，然后训练残差模型

# 迁移学习
model_ft = models.resnet50(pretrained=True)    
num_ftrs = model_ft.fc.in_features            
model_ft.fc = nn.Linear(num_ftrs, 12)           
model_ft = model_ft.to(device)         

criterion = nn.CrossEntropyLoss()       

optimizer_ft = torch.optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)   

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=5, gamma=0.1) 

model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler,num_epochs=15)    # 模型训练

4. 保存模型

没啥说的

def vali(M ,dataset):
    M.eval()
    with torch.no_grad():
        correct = 0
        for (data, target) in val_loader:
            data, target = data.to(device), target.to(device)

            pred = M(data)
            _, id = torch.max(pred, 1)
            correct += torch.sum(id == target.data)
        print("test accu: %.03f%%" % (100 * correct / len(dataset)))
    return (100 * correct / len(dataset)).item()
test_accu = int(vali(model_ft, img_datasets['val']) * 100)

model_name = 'val_{}.pkl'.format(test_accu)


torch.save(model_ft.state_dict(), os.path.join("./myModels", model_name))

5. 读取模型

# 加载模型
model_ft = models.resnet50(pretrained=True)    
num_ftrs = model_ft.fc.in_features            
model_ft.fc = nn.Linear(num_ftrs, 12)   # 注意更改维度         
model_ft = model_ft.to(device) # 注意要放入gpu，保持和参数数据类型一致

model_ft.load_state_dict(torch.load("./myModels/val_9343.pkl"))
vali(model_ft, img_datasets['val'])

输出结果：

test accu: 93.433%