29 lines
855 B
Python
29 lines
855 B
Python
import os
|
|
import logging
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from utils.model 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):
|
|
torch.save(self.model.state_dict(), os.path.join(save_dir, "model_state_dict.pt"))
|
|
|
|
def from_pretrained(self, load_dir):
|
|
state_dict = torch.load(load_dir, 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 |