Models Integrated

Author: AnirudhAchal

Net/extractors.py

@@ -162,13 +162,13 @@ def forward(self, x):
 class _DenseLayer(nn.Sequential):
     def __init__(self, num_input_features, growth_rate, bn_size, drop_rate):
         super(_DenseLayer, self).__init__()
-        self.add_module('norm.1', nn.BatchNorm2d(num_input_features)),
-        self.add_module('relu.1', nn.ReLU(inplace=True)),
-        self.add_module('conv.1', nn.Conv2d(num_input_features, bn_size *
+        self.add_module('norm1', nn.BatchNorm2d(num_input_features)),
+        self.add_module('relu1', nn.ReLU(inplace=True)),
+        self.add_module('conv1', nn.Conv2d(num_input_features, bn_size *
                                             growth_rate, kernel_size=1, stride=1, bias=False)),
-        self.add_module('norm.2', nn.BatchNorm2d(bn_size * growth_rate)),
-        self.add_module('relu.2', nn.ReLU(inplace=True)),
-        self.add_module('conv.2', nn.Conv2d(bn_size * growth_rate, growth_rate,
+        self.add_module('norm2', nn.BatchNorm2d(bn_size * growth_rate)),
+        self.add_module('relu2', nn.ReLU(inplace=True)),
+        self.add_module('conv2', nn.Conv2d(bn_size * growth_rate, growth_rate,
                                             kernel_size=3, stride=1, padding=1, bias=False)),
         self.drop_rate = drop_rate
 
@@ -332,33 +332,38 @@ def densenet(pretrained=True):
 def resnet18(pretrained=True):
     model = ResNet(BasicBlock, [2, 2, 2, 2])
     if pretrained:
-        load_weights_sequential(model, model_zoo.load_url(model_urls['resnet18']))
+        pass
+        # load_weights_sequential(model, model_zoo.load_url(model_urls['resnet18']))
     return model
 
 
 def resnet34(pretrained=True):
     model = ResNet(BasicBlock, [3, 4, 6, 3])
     if pretrained:
-        load_weights_sequential(model, model_zoo.load_url(model_urls['resnet34']))
+        pass
+        # load_weights_sequential(model, model_zoo.load_url(model_urls['resnet34']))
     return model
 
 
 def resnet50(pretrained=True):
     model = ResNet(Bottleneck, [3, 4, 6, 3])
     if pretrained:
-        load_weights_sequential(model, model_zoo.load_url(model_urls['resnet50']))
+        pass
+        # load_weights_sequential(model, model_zoo.load_url(model_urls['resnet50']))
     return model
 
 
 def resnet101(pretrained=True):
     model = ResNet(Bottleneck, [3, 4, 23, 3])
     if pretrained:
-        load_weights_sequential(model, model_zoo.load_url(model_urls['resnet101']))
+        pass
+        # load_weights_sequential(model, model_zoo.load_url(model_urls['resnet101']))
     return model
 
 
 def resnet152(pretrained=True):
     model = ResNet(Bottleneck, [3, 8, 36, 3])
     if pretrained:
-        load_weights_sequential(model, model_zoo.load_url(model_urls['resnet152']))
+        pass
+        # load_weights_sequential(model, model_zoo.load_url(model_urls['resnet152']))
     return model

Net/pspnet.py

@@ -40,7 +40,7 @@ def forward(self, x):
 
 
 class PSPNet(nn.Module):
-    def __init__(self, n_classes=18, sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet34',
+    def __init__(self, n_classes=20, sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet34',
                  pretrained=True):
         super().__init__()
         self.feats = getattr(extractors, backend)(pretrained)

train.py

@@ -16,13 +16,13 @@
 
 # Models
 models = {
-    'squeezenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='squeezenet'),
-    'densenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=1024, deep_features_size=512, backend='densenet'),
-    'resnet18': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet18'),
-    'resnet34': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet34'),
-    'resnet50': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet50'),
-    'resnet101': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet101'),
-    'resnet152': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet152')
+    'squeezenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='squeezenet', n_classes=20),
+    'densenet': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=1024, deep_features_size=512, backend='densenet', n_classes=20),
+    'resnet18': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet18', n_classes=20),
+    'resnet34': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet34', n_classes=20),
+    'resnet50': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet50', n_classes=20),
+    'resnet101': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet101', n_classes=20),
+    'resnet152': lambda: PSPNet(sizes=(1, 2, 3, 6), psp_size=2048, deep_features_size=1024, backend='resnet152', n_classes=20)
 }
 
 
@@ -64,13 +64,16 @@ def parse_arguments():
     parser = argparse.ArgumentParser(description='Human Parsing')
 
     # Add more arguments based on requirements later
-    parser.add_argument('-e', '--epochs', help='Set number of train epochs', default=100, type=int)
+    parser.add_argument('-e', '--epochs', help='Set number of train epochs', default=30, type=int)
     parser.add_argument('-b', '--batch-size', help='Set size of the batch', default=32, type=int)
     parser.add_argument('-d', '--data-path', help='Set path of dataset', default='.', type=str)
     parser.add_argument('-n', '--num-class', help='Set number of segmentation classes', default=20, type=int)
     parser.add_argument('-be', '--backend', help='Set Feature extractor', default='densenet', type=str)
     parser.add_argument('-s', '--snapshot', help='Set path to pre-trained weights', default=None, type=str)
-    parser.add_argument('-g', '--gpu', help='Set gpu [True / False]', default=False, type=bool)
+    parser.add_argument('-g', '--gpu', help='Set gpu [True / False]', default=False, action='store_true')
+    parser.add_argument('-lr', '--start-lr', help='Set starting learning rate', default=0.001, type=float)
+    parser.add_argument('-a', '--alpha', help='Set coefficient for classification loss term', default=1.0, type=float)
+    parser.add_argument('-m', '--milestones', type=str, default='10,20,30', help='Milestones for LR decreasing')
 
     # Mutually Exclusive Group 1 (Train / Eval)
     train_eval_parser = parser.add_mutually_exclusive_group(required=False)
@@ -121,13 +124,44 @@ def run_trained_model(model_ft, train_loader):
     models_path = os.path.join('./checkpoints', args.backend)
     os.makedirs(models_path, exist_ok=True)
 
-    train_loader = get_dataloader(args.data_path, train=args.train, batch_size=args.batch_size, num_class=args.num_class)
-
-    # Debug
-    for data in train_loader:
-        x, y, cls = data
-        break
-
-    # run_trained_model(models['torch_resnet50'], train_loader)
+    net, starting_epoch = build_network(args.snapshot, args.backend)
+    optimizer = optim.Adam(net.parameters(), lr=args.start_lr)
+    scheduler = MultiStepLR(optimizer, milestones=[int(x) for x in args.milestones.split(',')])
 
+    train_loader = get_dataloader(args.data_path, train=args.train, batch_size=args.batch_size, num_class=args.num_class)
 
+    for epoch in range(1+starting_epoch, 1+starting_epoch+args.epochs):
+        seg_criterion = nn.NLLLoss(weight=None)
+        cls_criterion = nn.BCEWithLogitsLoss(weight=None)
+        epoch_losses = []
+        net.train()
+
+        for count, (img, gt, gt_cls) in enumerate(train_loader):
+            # Input data
+            if args.gpu:
+                img, gt, gt_cls = img.cuda(), gt.cuda(), gt_cls.cuda()
+
+            img, gt, gt_cls = img, gt.long(), gt_cls.float()
+
+            # Forward pass
+            out, out_cls = net(img)
+            seg_loss, cls_loss = seg_criterion(out, gt), cls_criterion(out_cls, gt_cls)
+            loss = seg_loss + args.alpha * cls_loss
+
+            # Backward
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+
+            # Log
+            epoch_losses.append(loss.item())
+            status = '[{0}] step = {1}/{2}, loss = {3:0.4f} avg = {4:0.4f}, LR = {5:0.7f}'.format(
+                epoch, count, len(train_loader),
+                loss.item(), np.mean(epoch_losses), scheduler.get_lr()[0])
+            print(status)
+
+        scheduler.step()
+        if epoch % 10 == 0:
+            torch.save(net.state_dict(), os.path.join(models_path, '_'.join(["PSPNet", str(epoch)])))
+    
+    torch.save(net.state_dict(), os.path.join(models_path, '_'.join(["PSPNet", 'last'])))