xdecoder/Instruct-X-Decoder

1# --------------------------------------------------------2# X-Decoder -- Generalized Decoding for Pixel, Image, and Language3# Copyright (c) 2022 Microsoft4# Licensed under The MIT License [see LICENSE for details]5# Written by Xueyan Zou (xueyan@cs.wisc.edu)6# --------------------------------------------------------7 8import os9import logging10 11import torch12import torch.nn as nn13 14from utils.model_loading import align_and_update_state_dicts15 16logger = logging.getLogger(__name__)17 18 19class BaseModel(nn.Module):20    def __init__(self, opt, module: nn.Module):21        super(BaseModel, self).__init__()22        self.opt = opt23        self.model = module24 25    def forward(self, *inputs, **kwargs):26        outputs = self.model(*inputs, **kwargs)27        return outputs28 29    def save_pretrained(self, save_dir):30        save_path = os.path.join(save_dir, 'model_state_dict.pt')31        torch.save(self.model.state_dict(), save_path)32 33    def from_pretrained(self, load_path):34        state_dict = torch.load(load_path, map_location=self.opt['device'])35        state_dict = align_and_update_state_dicts(self.model.state_dict(), state_dict)36        self.model.load_state_dict(state_dict, strict=False)37        return self