x-decoder-video / xdecoder /BaseModel.py
fffiloni's picture
Create BaseModel.py
f6b8ea3
raw
history blame
1.22 kB
# --------------------------------------------------------
# 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