File size: 917 Bytes
5611f26 | 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 | import random
import torch
import numpy as np
def set_seed(seed):
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
def load_pretrain_model(net, weights):
net_keys = list(net.state_dict().keys())
weights_keys = list(weights.keys())
# assert(len(net_keys) <= len(weights_keys))
i = 0
j = 0
while i < len(net_keys) and j < len(weights_keys):
name_i = net_keys[i]
name_j = weights_keys[j]
if net.state_dict()[name_i].shape == weights[name_j].shape:
net.state_dict()[name_i].copy_(weights[name_j].cpu())
i += 1
j += 1
else:
i += 1
# print i, len(net_keys), j, len(weights_keys)
return net
|