def predict():
model = Resnet()
# original saved file with DataParallel
state_dict = torch.load('sgdr_rgb.pkl')
# create new OrderedDict that does not contain `module.`
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`
new_state_dict[name] = v
# load params
model.load_state_dict(new_state_dict)
#model.load_state_dict(torch.load('best.pkl'))
th = choose_th(model)
handle(model, th)
def main():
model = Resnet().to(device)
# trained_model = resnet18(pretrained=True)
# print(trained_model)
# model = nn.Sequential(*list(trained_model.children())[:-1],#[b, 512, 1, 1]
# Flatten(), # [b, 512 ,1,1] >=[b, 512]
# nn.Linear(512, 10),
# ).to(device)
print(model)
optimizer = optim.Adam(model.parameters(), lr=lr)
criteon = nn.CrossEntropyLoss()
best_acc, best_epoch = 0, 0
global_step = 0
viz.line([0], [0],
win='loss',
opts=dict(title='loss', xlabel='batch', ylabel='loss'))
viz.line([0], [0],
win='val_acc',
opts=dict(title='val_acc', xlabel='batch', ylabel='accuracy'))
for epoch in range(epochs):
for step, (x, y) in enumerate(train_loader):
# x: [b, 3, 224, 224], y: [b]
# print(x.shape,x,y.shape,y)
x, y = x.to(device), y.to(device)
model.train()
logits = model(x)
# print('logits is:', logits.cpu().detach().numpy())
loss = criteon(logits, y)
# print("loss:", loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()
viz.line([loss.item()], [global_step], win='loss', update='append')
global_step += 1
print('已进行到：', batchsz * global_step)
if epoch % 1 == 0:
val_acc = evalute(model, val_loader)
print(val_acc)
# viz.line([val_acc], [global_step], win='val_acc', update='append')
if val_acc > best_acc:
best_epoch = epoch
best_acc = val_acc
torch.save(model.state_dict(), 'best_canny.mdl')
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_canny.mdl'))
print('loaded from ckpt!')
test_acc = evalute(model, test_loader)
print('test acc:', test_acc)
