[Update] classmethod, from_config

Author: drawcodeboy

datasets/__init__.py

@@ -1,10 +1,8 @@
 from .mnist_dataset import MNIST_Dataset
 
-def load_dataset(**cfg):
+def load_dataset(cfg):
     if cfg['dataset'] == 'MNIST':
-        return MNIST_Dataset(root=cfg['root'],
-                             download=cfg['download'],
-                             mode=cfg['mode'])
+        return MNIST_Dataset.from_config(cfg)
 
     else:
         raise Exception(f"Dataset: {cfg['dataset']} is not supported.")

datasets/mnist_dataset.py

@@ -37,4 +37,10 @@ def __getitem__(self, idx):
 
         label = torch.tensor(label, dtype=torch.int64)
 
-        return image, label
+        return image, label
+    
+    @classmethod
+    def from_config(cls, cfg):
+        return cls(root=cfg['root'],
+                   download=cfg['download'],
+                   mode=cfg['mode'])

models/ConvNet/convnet.py

@@ -33,4 +33,10 @@ def forward(self, x):
 
         x = self.li(x)
 
-        return x
+        return x
+    
+    @classmethod
+    def from_config(cls, cfg):
+        return cls(in_channels=cfg['in_channels'],
+                   layers=cfg['layers'],
+                   class_num=cfg['class_num'])

models/ConvNet2/convnet2.py

@@ -33,4 +33,10 @@ def forward(self, x):
 
         x = self.li(x)
 
-        return x
+        return x
+    
+    @classmethod
+    def from_config(cls, cfg):
+        return cls(in_channels=cfg['in_channels'],
+                   layers=cfg['layers'],
+                   class_num=cfg['class_num'])

models/__init__.py

@@ -1,16 +1,12 @@
 from .ConvNet.convnet import ConvNet
 from .ConvNet2.convnet2 import ConvNet2
 
-def load_model(**cfg):
+def load_model(cfg):
     if cfg['name'] == 'ConvNet':
-        return ConvNet(in_channels=cfg['in_channels'],
-                       layers=cfg['layers'],
-                       class_num=cfg['class_num'])
+        return ConvNet.from_config(cfg)
 
     elif cfg['name'] == 'ConvNet2':
-        return ConvNet2(in_channels=cfg['in_channels'],
-                        layers=cfg['layers'],
-                        class_num=cfg['class_num'])
+        return ConvNet2.from_config(cfg)
 
     else:
         raise Exception(f"Model: {cfg['name']} is not supported.")