pytorch_proj/pytorch_tutorials/model_loading.py at master · sh2439/pytorch_proj · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
### Save and load the entire model

# Save the model
torch.save(model, PATH)

# Load the model
model = torch.load(PATH)
model.eval()


### Saving and loading a checkpoint

# Save the model/checkpoint
torch.save({'epoch':epoch,
			'model_state_dict':model.state_dict(),
			'optimizer_state_dict':optimizer.state_dict(),
			'loss':loss},
			PATH)

# Load and use the checkpoint
checkpoint = torch.load(PATH)

model = ModelClass()
optimizer = OptimizerClass()

model.load_state_dict.(checkpoint[model_state_dict])
optimizer.load_state_dict(checkpoint[optimizer_state_dict])
epoch = checkpoint['epoch']
loss = checkpoint['loss']


### Save and load model for Transfer Learning/Finetuning
torch.save(modelA.state_dict(),PATH)

modelB = ModelClass()
modelB.load_state_dict(torch.load(PATH), strict = False)


### Saving and loading models across devices

### Save on GPU, load on CPU
torch.save(model.state_dict(),PATH)

device = torch.device('cpu')
model = ModelClass()
model.load_state_dict(torch.load(PATH, map_location = device))

### Save on GPU, load on GPU
torch.save(model.state_dict(),PATH)

device = torch.device('cuda')
model = ModelClass()
model.load_state_dict(torch.load(PATH))
model.to(device)

### Save on CPU, load onn GPU
torch.save(model.state_dict(), PATH)

device = torch.device('cuda')
model = ModelClass()
model.load_state_dict(torch.load(PATH, map_location = 'cuda:0'))
model.to(device)