| 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 | |