| 10 | |
| 11 | |
| 12 | class BaseModel(nn.Module): |
| 13 | def __init__(self, opt, module: nn.Module): |
| 14 | super(BaseModel, self).__init__() |
| 15 | self.opt = opt |
| 16 | self.model = module |
| 17 | |
| 18 | def forward(self, *inputs, **kwargs): |
| 19 | outputs = self.model(*inputs, **kwargs) |
| 20 | return outputs |
| 21 | |
| 22 | def save_pretrained(self, save_path): |
| 23 | torch.save(self.model.state_dict(), save_path) |
| 24 | |
| 25 | def from_pretrained(self, load_dir): |
| 26 | state_dict = torch.load(load_dir, map_location='cpu') |
| 27 | if 'model' in state_dict: |
| 28 | state_dict=state_dict['model'] |
| 29 | state_dict={k[6:]:v for k,v in state_dict.items() if k.startswith('model.')} |
| 30 | # for k in self.model.state_dict(): |
| 31 | # if k not in state_dict: |
| 32 | # assert k[:-2] in state_dict |
| 33 | # state_dict[k]=state_dict.pop(k[:-2]) |
| 34 | state_dict = align_and_update_state_dicts(self.model.state_dict(), state_dict) |
| 35 | self.model.load_state_dict(state_dict, strict=False) |
| 36 | return self |
no outgoing calls
no test coverage detected