feat: add cli support for switching model

detector/model.py

@@ -11,9 +11,10 @@ import pytorch_lightning as ptl
 
 
 class ResNet18Regressor(nn.Module):
-    def __init__(self, regression_use_tanh: bool = False):
+    def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False):
         super().__init__()
-        self.model = torchvision.models.resnet18(weights=False)
+        weights = torchvision.models.ResNet18_Weights.DEFAULT if pretrained else None
+        self.model = torchvision.models.resnet18(weights=weights)
         self.model.fc = nn.Linear(512, config.FONT_COUNT + 12)
         self.regression_use_tanh = regression_use_tanh
 
@@ -28,9 +29,10 @@ class ResNet18Regressor(nn.Module):
 
 
 class ResNet34Regressor(nn.Module):
-    def __init__(self, regression_use_tanh: bool = False):
+    def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False):
         super().__init__()
-        self.model = torchvision.models.resnet34(weights=False)
+        weights = torchvision.models.ResNet34_Weights.DEFAULT if pretrained else None
+        self.model = torchvision.models.resnet34(weights=weights)
         self.model.fc = nn.Linear(512, config.FONT_COUNT + 12)
         self.regression_use_tanh = regression_use_tanh
 
@@ -45,9 +47,10 @@ class ResNet34Regressor(nn.Module):
 
 
 class ResNet50Regressor(nn.Module):
-    def __init__(self, regression_use_tanh: bool = False):
+    def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False):
         super().__init__()
-        self.model = torchvision.models.resnet50(weights=False)
+        weights = torchvision.models.ResNet50_Weights.DEFAULT if pretrained else None
+        self.model = torchvision.models.resnet50(weights=weights)
         self.model.fc = nn.Linear(2048, config.FONT_COUNT + 12)
         self.regression_use_tanh = regression_use_tanh
 
@@ -62,9 +65,10 @@ class ResNet50Regressor(nn.Module):
 
 
 class ResNet101Regressor(nn.Module):
-    def __init__(self, regression_use_tanh: bool = False):
+    def __init__(self, pretrained: bool = False, regression_use_tanh: bool = False):
         super().__init__()
-        self.model = torchvision.models.resnet101(weights=False)
+        weights = torchvision.models.ResNet101_Weights.DEFAULT if pretrained else None
+        self.model = torchvision.models.resnet101(weights=weights)
         self.model.fc = nn.Linear(2048, config.FONT_COUNT + 12)
         self.regression_use_tanh = regression_use_tanh
 

train.py

@@ -12,9 +12,42 @@ from utils import get_current_tag
 torch.set_float32_matmul_precision("high")
 
 parser = argparse.ArgumentParser()
-parser.add_argument("-d", "--devices", nargs="*", type=int, default=[0])
-parser.add_argument("-b", "--single-batch-size", type=int, default=64)
-parser.add_argument("-c", "--checkpoint", type=str, default=None)
+parser.add_argument(
+    "-d",
+    "--devices",
+    nargs="*",
+    type=int,
+    default=[0],
+    help="GPU devices to use (default: [0])",
+)
+parser.add_argument(
+    "-b",
+    "--single-batch-size",
+    type=int,
+    default=64,
+    help="Batch size of single device (default: 64)",
+)
+parser.add_argument(
+    "-c",
+    "--checkpoint",
+    type=str,
+    default=None,
+    help="Trainer checkpoint path (default: None)",
+)
+parser.add_argument(
+    "-m",
+    "--model",
+    type=str,
+    default="resnet18",
+    choices=["resnet18", "resnet34", "resnet50", "resnet101"],
+    help="Model to use (default: resnet18)",
+)
+parser.add_argument(
+    "-p",
+    "--pretrained",
+    action="store_true",
+    help="Use pretrained model for ResNet (default: False)",
+)
 
 args = parser.parse_args()
 
@@ -76,7 +109,24 @@ trainer = ptl.Trainer(
     deterministic=True,
 )
 
-model = ResNet50Regressor(regression_use_tanh=regression_use_tanh)
+if args.model == "resnet18":
+    model = ResNet18Regressor(
+        pretrained=args.pretrained, regression_use_tanh=regression_use_tanh
+    )
+elif args.model == "resnet34":
+    model = ResNet34Regressor(
+        pretrained=args.pretrained, regression_use_tanh=regression_use_tanh
+    )
+elif args.model == "resnet50":
+    model = ResNet50Regressor(
+        pretrained=args.pretrained, regression_use_tanh=regression_use_tanh
+    )
+elif args.model == "resnet101":
+    model = ResNet101Regressor(
+        pretrained=args.pretrained, regression_use_tanh=regression_use_tanh
+    )
+else:
+    raise NotImplementedError()
 
 detector = FontDetector(
     model=model,