Create BaseModel.py

xdecoder/BaseModel.py

@@ -0,0 +1,37 @@
+# --------------------------------------------------------
+# 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 ([email protected])
+# --------------------------------------------------------
+
+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