1 import torch
2 from torch import optim,nn
3 import visdom
4 import torchvision
5 from torch.utils.data import DataLoader
6
7 from pokemon import Pokemon
8
9 # from resnet import ResNet18
10 # 可以加载直接加载好的状态
11 from torchvision.models import resnet18
12
13 from utils import Flatten
14
15 batchsz = 32
16 lr = 1e-3
17 epochs = 10
18
19 device = torch.device('cuda')
20 # 设置随机种子保证能够复现出来
21 torch.manual_seed(1234)
22
23 train_db = Pokemon('pokemon',224,mode = 'train')
24 val_db = Pokemon('pokemon',224,mode = 'val')
25 test_db = Pokemon('pokemon',224,mode = 'test')
26
27 train_loader = DataLoader(train_db,batch_size = batchsz,shuffle = True,num_workers = 4)
28 val_loader = DataLoader(val_db,batch_size = batchsz,num_workers = 2)
29 test_loader = DataLoader(test_db,batch_size = batchsz,num_workers = 2)
30
31 # visdom
32 viz = visdom.Visdom()
33
34 def evalute(model,loader):
35
36 correct = 0
37 total = len(loader.dataset)
38
39 for x,y in loader:
40 x,y = x.to(device),y.to(device)
41 with torch.no_grad():
42 logits = model(x)
43 pred = logits.argmax(dim = 1)
44 correct += torch.eq(pred,y).sum().float().item()
45
46 return correct / total
47
48 def main():
49
50 # model = ResNet18(5).to(device)
51 trained_model = resnet18(pretrained = True)
52 # 取出前17层,加*打散数据
53 model = nn.Sequential(*list(train_model.children())[:-1], # [b,512,1,1]
54 Flatten(), # [b,512,1,1] --> [b,512]
55 nn.Linear(512,5)
56 ).to(device)
57
58 optimizer = optim.Adam(model.parameters().lr = lr)
59 criteon = nn.CrossEntropyLoss
60
61 best_acc,best_epoch = 0,0
62 global_step = 0
63 # visdom
64 viz.line([0],[-1],win = 'loss',opts = dict(title = 'loss'))
65 viz.line([0],[-1],win = 'val_acc',opts = dict(title = 'val_acc'))
66
67 for epoch in range(epochs):
68
69 for step,(x,y) in enumerate(train_loader):
70
71 # x: [b,3,224,224] ,y : [b]
72 x,y = x.to(device),y.to(device)
73
74 # logits是没经过loss的
75 logits = model(x)
76 # CrossEntropyLoss会在内部进行onehot，所以不需要自己写
77 loss = criteon(logits,y).item()
78
79 optimizer.zero_grad()
80 loss.backward()
81 optimizer.step()
82
83 # visdom
84 viz.line([loss.item()],[global_step],win = 'loss',update = 'append')
85 global_step += 1
86
87 if epoch % 2 == 0:
88
89 val_acc = evalute(model,val_loader)
90
91 if val_acc > best_acc:
92 best_epoch = epoch
93 best_acc = val_acc
94
95 torch.save(model.state_dict(),'best.mdl')
96 # visdom
97 viz.line([val_acc],[global_step],win = 'val_acc',update = 'append')
98
99 print('best_acc:',best_acc,'best_epoch',best_epoch)
100
101 model.load_state_dict(torch.load('best.mdl'))
102 print('loaded from skpt!')
103
104 test_acc = evalute(model,test_loader)
105 print('test_acc',test_acc)
106
107
108 if __name__ == '__main__'
109 main()
