Spaces:
Build error
Build error
| # -------------------------------------------------------- | |
| # X-Decoder -- Generalized Decoding for Pixel, Image, and Language | |
| # Copyright (c) 2022 Microsoft | |
| # Licensed under The MIT License [see LICENSE for details] | |
| # Written by Xueyan Zou (xueyan@cs.wisc.edu) | |
| # -------------------------------------------------------- | |
| import os | |
| import logging | |
| import torch | |
| import torch.nn as nn | |
| from utils.model_loading import align_and_update_state_dicts | |
| logger = logging.getLogger(__name__) | |
| class BaseModel(nn.Module): | |
| def __init__(self, opt, module: nn.Module): | |
| super(BaseModel, self).__init__() | |
| self.opt = opt | |
| self.model = module | |
| def forward(self, *inputs, **kwargs): | |
| outputs = self.model(*inputs, **kwargs) | |
| return outputs | |
| def save_pretrained(self, save_dir): | |
| save_path = os.path.join(save_dir, 'model_state_dict.pt') | |
| torch.save(self.model.state_dict(), save_path) | |
| def from_pretrained(self, load_path): | |
| state_dict = torch.load(load_path, map_location=self.opt['device']) | |
| state_dict = align_and_update_state_dicts(self.model.state_dict(), state_dict) | |
| self.model.load_state_dict(state_dict, strict=False) | |
| return self |