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)