龙良曲pytorch学习笔记_迁移学习

import torch
from torch import optim,nn
import visdom
import torchvision
from torch.utils.data import DataLoader

from pokemon import Pokemon

# from resnet import ResNet18
# 可以加载直接加载好的状态
from torchvision.models import resnet18

from utils import Flatten

batchsz = 32
lr = 1e-3
epochs = 10

device = torch.device('cuda')
# 设置随机种子保证能够复现出来
torch.manual_seed(1234)

train_db = Pokemon('pokemon',224,mode = 'train')
val_db = Pokemon('pokemon',224,mode = 'val')
test_db = Pokemon('pokemon',224,mode = 'test')

train_loader = DataLoader(train_db,batch_size = batchsz,shuffle = True,num_workers = 4)
val_loader = DataLoader(val_db,batch_size = batchsz,num_workers = 2)
test_loader = DataLoader(test_db,batch_size = batchsz,num_workers = 2)

# visdom
viz = visdom.Visdom()

def evalute(model,loader):
    
    correct = 0
    total = len(loader.dataset)
    
    for x,y in loader:
        x,y = x.to(device),y.to(device)
        with torch.no_grad():
            logits = model(x)
            pred = logits.argmax(dim = 1)
            correct += torch.eq(pred,y).sum().float().item()
        
    return correct / total

def main():
    
    # model = ResNet18(5).to(device)
    trained_model = resnet18(pretrained = True)
    # 取出前17层,加*打散数据
    model = nn.Sequential(*list(train_model.children())[:-1], # [b,512,1,1]
                          Flatten(), # [b,512,1,1] --> [b,512]
                          nn.Linear(512,5)
                          ).to(device)
    
    optimizer = optim.Adam(model.parameters().lr = lr)
    criteon = nn.CrossEntropyLoss
    
    best_acc,best_epoch = 0,0
    global_step = 0
    # visdom
    viz.line([0],[-1],win = 'loss',opts = dict(title = 'loss'))
    viz.line([0],[-1],win = 'val_acc',opts = dict(title = 'val_acc'))
    
    for epoch in range(epochs):
        
        for step,(x,y) in enumerate(train_loader):
            
            # x: [b,3,224,224] ,y : [b]
            x,y = x.to(device),y.to(device)
            
            # logits是没经过loss的
            logits = model(x)
            # CrossEntropyLoss会在内部进行onehot，所以不需要自己写
            loss = criteon(logits,y).item()
            
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            # visdom
            viz.line([loss.item()],[global_step],win = 'loss',update = 'append')
            global_step += 1
            
        if epoch % 2 == 0:
        
            val_acc = evalute(model,val_loader)
            
            if val_acc > best_acc:
                best_epoch = epoch
                best_acc = val_acc
                
                torch.save(model.state_dict(),'best.mdl')
                # visdom
                viz.line([val_acc],[global_step],win = 'val_acc',update = 'append')
                
    print('best_acc:',best_acc,'best_epoch',best_epoch)
    
    model.load_state_dict(torch.load('best.mdl'))
    print('loaded from skpt!')
    
    test_acc = evalute(model,test_loader)
    print('test_acc',test_acc)


if __name__ == '__main__'
    main()