Fix config error

Author: synml

eval.py

@@ -99,7 +99,7 @@ def evaluate(model, testloader, num_classes: int, device):
     model.eval()
 
     # 모델 평가
-    val_loss, iou, miou, fps = evaluate(model, testloader, config['model']['num_classes'], device)
+    val_loss, iou, miou, fps = evaluate(model, testloader, config[config['model']]['num_classes'], device)
 
     # 평가 결과를 csv 파일로 저장
     os.makedirs('result', exist_ok=True)

train.py

@@ -23,7 +23,7 @@
 
     # 3. Loss function, optimizer, lr scheduler
     criterion = nn.CrossEntropyLoss()
-    optimizer = torch.optim.Adam(model.parameters(), lr=config['model']['lr'])
+    optimizer = torch.optim.Adam(model.parameters(), lr=config[config['model']]['lr'])
     scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, min_lr=0.0001)
 
     # 4. Tensorboard
@@ -34,7 +34,7 @@
     log_loss = tqdm.tqdm(total=0, position=2, bar_format='{desc}', leave=False)
     prev_miou = 0.0
     prev_val_loss = 0.0
-    for epoch in tqdm.tqdm(range(config['model']['epoch']), desc='Epoch'):
+    for epoch in tqdm.tqdm(range(config[config['model']]['epoch']), desc='Epoch'):
         model.train()
 
         for batch_idx, (images, masks) in enumerate(tqdm.tqdm(trainloader, desc='Train', leave=False)):
@@ -58,7 +58,7 @@
             writer.add_scalar('Train Loss', loss.item(), len(trainloader) * epoch + batch_idx)
 
         # 모델 평가
-        val_loss, _, miou, _ = eval.evaluate(model, testloader, config['model']['num_classes'], device)
+        val_loss, _, miou, _ = eval.evaluate(model, testloader, config[config['model']]['num_classes'], device)
         writer.add_scalar('Validation Loss', val_loss, epoch)
         writer.add_scalar('mIoU', miou, epoch)
 

utils.py

@@ -30,17 +30,17 @@ def get_model(config: dict, pretrained=False) -> torch.nn.Module:
     assert isinstance(pretrained, bool)
 
     if config['model'] == 'UNet':
-        model = models.unet.UNet(config['model']['num_classes'])
+        model = models.unet.UNet(config[config['model']]['num_classes'])
     elif config['model'] == 'Proposed':
-        model = models.proposed.Proposed(config['model']['num_classes'])
+        model = models.proposed.Proposed(config[config['model']]['num_classes'])
     elif config['model'] == 'Backbone':
-        model = models.backbone.Backbone(config['model']['num_classes'])
+        model = models.backbone.Backbone(config[config['model']]['num_classes'])
     else:
         raise NameError('Wrong model_name.')
 
     if pretrained:
-        if os.path.exists(config['model']['pretrained_weights']):
-            model.load_state_dict(torch.load(config['model']['pretrained_weights']))
+        if os.path.exists(config[config['model']]['pretrained_weights']):
+            model.load_state_dict(torch.load(config[config['model']]['pretrained_weights']))
         else:
             print('FileNotFound: pretrained_weights (' + config['model'] + ')')
     return model
@@ -77,7 +77,7 @@ def set_cityscapes(self):
                                                    target_type='semantic',
                                                    transforms=self.transforms)
         trainloader = torch.utils.data.DataLoader(trainset,
-                                                  batch_size=self.config['model']['batch_size'],
+                                                  batch_size=self.config[self.config['model']]['batch_size'],
                                                   shuffle=True,
                                                   num_workers=self.config['dataset']['num_workers'],
                                                   pin_memory=True)
@@ -88,7 +88,7 @@ def set_cityscapes(self):
                                                   transform=self.transform,
                                                   target_transform=self.target_transform)
         testloader = torch.utils.data.DataLoader(testset,
-                                                 batch_size=self.config['model']['batch_size'],
+                                                 batch_size=self.config[self.config['model']]['batch_size'],
                                                  shuffle=False,
                                                  num_workers=self.config['dataset']['num_workers'])