(self, path)
| 359 | self.best_model = deepcopy(self.model) |
| 360 | |
| 361 | def load_pretrained(self, path): |
| 362 | |
| 363 | with open(os.path.join(path, 'config.pkl'), 'rb') as f: |
| 364 | config = pickle.load(f) |
| 365 | |
| 366 | del config['device'], config['num_genes'], config['num_perts'] |
| 367 | self.model_initialize(**config) |
| 368 | self.config = config |
| 369 | |
| 370 | state_dict = torch.load(os.path.join(path, 'model.pt'), map_location = torch.device('cpu')) |
| 371 | if next(iter(state_dict))[:7] == 'module.': |
| 372 | # the pretrained model is from data-parallel module |
| 373 | from collections import OrderedDict |
| 374 | new_state_dict = OrderedDict() |
| 375 | for k, v in state_dict.items(): |
| 376 | name = k[7:] # remove `module.` |
| 377 | new_state_dict[name] = v |
| 378 | state_dict = new_state_dict |
| 379 | |
| 380 | self.model.load_state_dict(state_dict) |
| 381 | self.model = self.model.to(self.device) |
| 382 | self.best_model = self.model |
| 383 | |
| 384 | def save_model(self, path): |
| 385 | if not os.path.exists(path): |
no test coverage detected