save、load
一、load
1、cpu——cpu、gpu——gpu
checkpoint = t.load(_.pth)
2、cpu——gpu
checkpoint = t.load(_.pth, map_location=lambda storage, loc: storage.cuda(i))
3、gpu——cpu
checkpoint = t.load(_.pth, map_location=lambda storage, loc: storage)
checkpoint = t.load(_.pth, map_location='cpu')
model.load_state_dict(checkpoint['state_dict'])
二、save
t.save({
'epoch' : epoch,
'state_dict' : model.state_dict()
},_.pth)