model: Remove input channels parameter

Author: synml

demo.py

@@ -18,7 +18,7 @@
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.utils.get_model(model_name, 3, config['num_classes'], config['pretrained_weights']).to(device)
+    model = utils.utils.get_model(model_name, config['num_classes'], config['pretrained_weights']).to(device)
 
     # 이미지 이름 불러오기
     image_names = []

eval.py

@@ -95,7 +95,7 @@ def evaluate(model, testloader, num_classes: int, device):
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.utils.get_model(model_name, 3, config['num_classes'], config['pretrained_weights']).to(device)
+    model = utils.utils.get_model(model_name, config['num_classes'], config['pretrained_weights']).to(device)
 
     # 모델 평가
     val_loss, iou, miou, fps = evaluate(model, testloader, config['num_classes'], device)

models/proposed.py

@@ -91,11 +91,11 @@ def forward(self, x):
 
 
 class Proposed(nn.Module):
-    def __init__(self, num_channels: int, num_classes: int):
+    def __init__(self, num_classes: int):
         super(Proposed, self).__init__()
         resnet50 = torchvision.models.resnet50(pretrained=True)
 
-        self.encode1 = self.double_conv(num_channels, 64)
+        self.encode1 = self.double_conv(3, 64)
         self.encode2 = resnet50.layer1  # 256
         self.encode3 = resnet50.layer2  # 512
         self.encode4 = resnet50.layer3  # 1024
@@ -154,7 +154,7 @@ def forward(self, x):
 
 if __name__ == '__main__':
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = Proposed(3, 8).to(device)
+    model = Proposed(8).to(device)
     model.eval()
 
     torchsummary.torchsummary.summary(model, (3, 256, 512))

models/unet.py

@@ -6,10 +6,10 @@
 
 
 class UNet(nn.Module):
-    def __init__(self, num_channels: int, num_classes: int):
+    def __init__(self, num_classes: int):
         super(UNet, self).__init__()
 
-        self.encode1 = self.double_conv(num_channels, 64)
+        self.encode1 = self.double_conv(3, 64)
         self.encode2 = self.double_conv(64, 128)
         self.encode3 = self.double_conv(128, 256)
         self.encode4 = self.double_conv(256, 512)
@@ -58,7 +58,7 @@ def forward(self, x):
 
 if __name__ == '__main__':
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = UNet(3, 8).to(device)
+    model = UNet(8).to(device)
     model.eval()
 
     torchsummary.torchsummary.summary(model, (3, 256, 512))

train.py

@@ -19,7 +19,7 @@
 
     # 2. Model
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-    model = utils.utils.get_model(model_name, 3, config['num_classes']).to(device)
+    model = utils.utils.get_model(model_name, config['num_classes']).to(device)
 
     # 3. Loss function, optimizer, lr scheduler
     criterion = nn.CrossEntropyLoss()

utils/utils.py

@@ -31,11 +31,11 @@ def load_config():
     return model_name, config
 
 
-def get_model(model_name: str, num_channels: int, num_classes: int, pretrained: str = None) -> torch.nn.Module:
+def get_model(model_name: str, num_classes: int, pretrained: str = None) -> torch.nn.Module:
     if model_name == 'UNet':
-        model = models.unet.UNet(num_channels, num_classes)
+        model = models.unet.UNet(num_classes)
     elif model_name == 'Proposed':
-        model = models.proposed.Proposed(num_channels, num_classes)
+        model = models.proposed.Proposed(num_classes)
     else:
         raise NameError('Wrong model_name.')