Fix Model out memory

configs/efficientdet-d0.yaml

@@ -80,7 +80,7 @@ CRITERION:
     neg_pos: 3
     neg_overlap: 0.5
     encode_target: False
-    use_gpu: True
+    use_gpu: False
 
 OPTIMIZER:
   PY: torch.optim

data/config.py

@@ -15,7 +15,7 @@
     'num_classes': 21,
     'lr_steps': (80000, 100000, 120000),
     'max_iter': 120000,
-    'feature_maps': [64, 32, 16, 8, 4],
+    'feature_maps': [40, 80, 112, 192, 320],
     'min_dim': 512,
     'steps': [8, 16, 32, 64, 128],
     'min_sizes': [30, 60, 111, 162, 213],

layers/modules/multibox_loss.py

@@ -76,13 +76,11 @@ def forward(self, predictions, targets):
             loc_t = loc_t.cuda()
             conf_t = conf_t.cuda()
 
-        print('loc_t: {}, conf_t: {}'.format(loc_t.is_cuda, conf_t.is_cuda))
+
         # wrap targets
         loc_t = Variable(loc_t, requires_grad=False)
         conf_t = Variable(conf_t, requires_grad=False)      
         pos = conf_t > 0
-
-        print('pos: ', pos.is_cuda)
         num_pos = pos.sum(dim=1, keepdim=True)
         # Localization Loss (Smooth L1)
         # Shape: [batch,num_priors,4]
@@ -91,7 +89,7 @@ def forward(self, predictions, targets):
         loc_p = loc_data[pos_idx].view(-1, 4)
         loc_t = loc_t[pos_idx].view(-1, 4)
         loss_l = F.smooth_l1_loss(loc_p, loc_t, size_average=False)
-
+        
         # Compute max conf across batch for hard negative mining
         batch_conf = conf_data.view(-1, self.num_classes)
         loss_c = log_sum_exp(batch_conf) - batch_conf.gather(1, conf_t.view(-1, 1))

models/bifpn.py

@@ -1,48 +1,81 @@
 import torch
-import torch.nn as nn 
+import torch.nn as nn
+
 
 class BiFPN(nn.Module):
     def __init__(self,
-                num_channels):
+                 num_channels):
         super(BiFPN, self).__init__()
         self.num_channels = num_channels
+        self.input_image = 512
+        self.top_down = [
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+        ]
+        self.down = [
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+            nn.Conv2d(in_channels=self.num_channels, out_channels=self.num_channels,
+                      kernel_size=3, padding=1, groups=self.num_channels),
+        ]
 
     def forward(self, inputs):
         num_channels = self.num_channels
         P3_in, P4_in, P5_in, P6_in, P7_in = inputs
 
-        P7_up = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P7_in)
+        P7_up = self.top_down[0](P7_in)
         scale = int(P6_in.size(3)/P7_up.size(3))
-        P6_up = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P6_in+self.Resize(scale_factor=scale)(P7_up))
+        P6_up = self.top_down[1](
+            P6_in + self.Resize(scale_factor=scale)(P7_up))
         scale = int(P5_in.size(3)/P6_up.size(3))
-        P5_up = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P5_in+self.Resize(scale_factor=scale)(P6_up))
+        P5_up = self.top_down[2](
+            P5_in + self.Resize(scale_factor=scale)(P6_up))
         scale = int(P4_in.size(3)/P5_up.size(3))
-        P4_up = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P4_in+self.Resize(scale_factor=scale)(P5_up))
+        P4_up = self.top_down[3](
+            P4_in + self.Resize(scale_factor=scale)(P5_up))
         scale = int(P3_in.size(3)/P4_up.size(3))
-        P3_out = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P3_in+self.Resize(scale_factor=scale)(P4_up))
+        P3_out = self.top_down[4](
+            P3_in + self.Resize(scale_factor=scale)(P4_up))
 
         kernel_size = int(P3_out.size(3)/P4_up.size(3))
-        P4_out = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P4_in + P4_up+nn.MaxPool2d(kernel_size=kernel_size)(P3_out))
+        P4_out = self.down[0](
+            P4_in+P4_up+nn.MaxPool2d(kernel_size=kernel_size)(P3_out))
         kernel_size = int(P4_out.size(3)/P5_up.size(3))
-        P5_out = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P5_in + P5_up+nn.MaxPool2d(kernel_size=kernel_size)(P4_out))
+        P5_out = self.down[1](
+            P5_in+P5_up+nn.MaxPool2d(kernel_size=kernel_size)(P4_out))
         kernel_size = int(P5_out.size(3)/P6_up.size(3))
-        P6_out = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P6_in + P6_up+nn.MaxPool2d(kernel_size=kernel_size)(P5_out))
+        P6_out = self.down[2](
+            P6_in+P6_up+nn.MaxPool2d(kernel_size=kernel_size)(P5_out))
         kernel_size = int(P6_out.size(3)/P7_up.size(3))
-        P7_out = self.Conv(in_channels=num_channels, out_channels=num_channels, kernel_size=1, stride=1, padding=0, groups=num_channels)(P7_in + P7_up+nn.MaxPool2d(kernel_size=kernel_size)(P6_out))
+        P7_out = self.down[3](
+            P7_in+P7_up+nn.MaxPool2d(kernel_size=kernel_size)(P6_out))
         return P3_out, P4_out, P5_out, P6_out, P7_out
 
-
-
     @staticmethod
-    def Conv(in_channels, out_channels, kernel_size, stride, padding, groups = 1):
+    def Conv(in_channels, out_channels, kernel_size, stride, padding, groups=1):
         features = nn.Sequential(
-            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=groups),
+            nn.Conv2d(in_channels=in_channels, out_channels=out_channels,
+                      kernel_size=kernel_size, stride=stride, padding=padding, groups=groups),
             nn.BatchNorm2d(num_features=out_channels),
             nn.ReLU()
         )
-        return features 
+        return features
+
     @staticmethod
     def Resize(scale_factor=2, mode='nearest'):
         upsample = nn.Upsample(scale_factor=scale_factor, mode=mode)
         return upsample
-

models/efficientdet.py

@@ -1,72 +1,95 @@
-import torch 
+import torch
 import torch.nn as nn
 from torch.autograd import Variable
 from models.efficientnet import EfficientNet
 from models.bifpn import BiFPN
 from layers.functions import PriorBox
 from data import voc, coco
+
+
+
 class EfficientDet(nn.Module):
     def __init__(self,
-                num_class = 21,
-                levels = 3,
-                num_channels = 128,
-                model_name = 'efficientnet-b0'):
+                 num_class=21,
+                 levels=3,
+                 num_channels=128,
+                 model_name='efficientnet-b0'):
         super(EfficientDet, self).__init__()
-        self.num_class = num_class 
+        self.num_class = num_class
         self.levels = levels
         self.num_channels = num_channels
         self.efficientnet = EfficientNet.from_pretrained(model_name)
-        self.bifpn = BiFPN(num_channels = self.num_channels)
+        self.bifpn = BiFPN(num_channels=self.num_channels)
         self.cfg = (coco, voc)[num_class == 21]
         self.priorbox = PriorBox(self.cfg)
         self.priors = Variable(self.priorbox.forward(), volatile=True)
-
-
+        self.num_anchor = 9
+
+        self.input_image = 512
+        self.Conv = [
+            nn.Conv2d(in_channels=40,
+                      out_channels=self.num_channels, kernel_size=3, padding=1, stride=1),
+            nn.Conv2d(in_channels=80,
+                      out_channels=self.num_channels, kernel_size=3, padding=1),
+            nn.Conv2d(in_channels=112,
+                      out_channels=self.num_channels, kernel_size=3, padding=1),
+            nn.Conv2d(in_channels=192,
+                      out_channels=self.num_channels, kernel_size=3, padding=1),
+            nn.Conv2d(in_channels=320,
+                      out_channels=self.num_channels, kernel_size=3, padding=1),
+        ]
+        self.class_module = list()
+        self.regress_module = list()
+        for _ in range(3, 8):
+            self.class_module.append(
+                nn.Sequential(
+                    nn.Conv2d(in_channels=self.num_channels, out_channels=64,
+                              kernel_size=2, stride=1),
+                    nn.Conv2d(in_channels=64, out_channels=self.num_anchor * num_class, kernel_size=2, stride=1)
+                )
+            )
+            self.regress_module.append(
+                nn.Sequential(
+                    nn.Conv2d(in_channels=self.num_channels, out_channels=64,
+                              kernel_size=2, stride=1),
+                    nn.Conv2d(
+                        in_channels=64, out_channels=self.num_anchor * 4, kernel_size=2, stride=1)
+                )
+            )
+
     def forward(self, inputs):
+
         P1, P2, P3, P4, P5, P6, P7 = self.efficientnet(inputs)
-        P3 = self.bifpn.Conv(in_channels=P3.size(1), out_channels=self.num_channels, kernel_size=1, stride=1, padding=0)(P3)
-        P4 = self.bifpn.Conv(in_channels=P4.size(1), out_channels=self.num_channels, kernel_size=1, stride=1, padding=0)(P4)
-        P5 = self.bifpn.Conv(in_channels=P5.size(1), out_channels=self.num_channels, kernel_size=1, stride=1, padding=0)(P5)
-        P6 = self.bifpn.Conv(in_channels=P6.size(1), out_channels=self.num_channels, kernel_size=1, stride=1, padding=0)(P6)
-        P7 = self.bifpn.Conv(in_channels=P7.size(1), out_channels=self.num_channels, kernel_size=1, stride=1, padding=0)(P7)
+        P3 = self.Conv[0](P3)
+        P4 = self.Conv[1](P4)
+        P5 = self.Conv[2](P5)
+        P6 = self.Conv[3](P6)
+        P7 = self.Conv[4](P7)
         for _ in range(self.levels):
             P3, P4, P5, P6, P7 = self.bifpn([P3, P4, P5, P6, P7])
-        P = [P3, P4, P5, P6, P7]
-
-        features_class = [self.class_net(p, self.num_class) for p in P]
-        features_class = torch.cat(features_class, axis=0)
-        features_bbox = [self.regression_net(p) for p in P]
-        features_bbox = torch.cat(features_bbox, axis=0)
+        feature_classes = []
+        feature_bboxes = []
+        for i, p in enumerate([P3, P4, P5, P6, P7]):
+            feature_class = self.class_module[i](p)
+            feature_class = feature_class.view(-1, self.num_class)
+            feature_class = nn.Sigmoid()(feature_class)
+            feature_classes.append(feature_class)
+
+            feature_bbox = self.regress_module[i](p)
+            feature_bbox = feature_bbox.view(-1, 4)
+            feature_bbox = nn.Sigmoid()(feature_bbox)
+            feature_bboxes.append(feature_bbox)
+        feature_classes = torch.cat(feature_classes, axis=0)
+        feature_bboxes = torch.cat(feature_bboxes, axis=0)
 
         output = (
-                features_bbox.view(inputs.size(0), -1, 4),
-                features_class.view(inputs.size(0), -1, self.num_class),
-                self.priors
-            )
+            feature_bboxes.view(inputs.size(0), -1, 4),
+            feature_classes.view(inputs.size(0), -1, self.num_class),
+            self.priors
+        )
         return output
-
-    @staticmethod
-    def class_net(features, num_class, num_anchor=3):
-        features = nn.Sequential(
-            nn.Conv2d(in_channels=features.size(1), out_channels=features.size(2), kernel_size = 2, stride=1),
-            nn.Conv2d(in_channels=features.size(2), out_channels=num_anchor*num_class, kernel_size = 2, stride=1)
-        )(features)
-        features = features.view(-1, num_class)
-        features = nn.Sigmoid()(features)
-        return features 
-
-    @staticmethod
-    def regression_net(features, num_anchor=3):
-        features = nn.Sequential(
-            nn.Conv2d(in_channels=features.size(1), out_channels=features.size(2), kernel_size = 2, stride=1),
-            nn.Conv2d(in_channels=features.size(2), out_channels=num_anchor*4, kernel_size = 2, stride=1)
-        )(features)
-        features = features.view(-1, 4)
-        features = nn.Sigmoid()(features)
-        return features 
-
 
-if __name__ =='__main__':
-    inputs = torch.randn(4, 3, 640, 640)
-    model = EfficientDet(levels=10)
-    output = model(inputs)
+if __name__ == '__main__':
+    inputs = torch.randn(4, 3, 512, 512)
+    model = EfficientDet(levels=3)
+    output = model(inputs)

test.py

@@ -0,0 +1,10 @@
+
+
+def sum():
+    a = 10
+    b = 20
+    for name in dir():
+        if not name.startswith('_'):
+            del locals()[name]
+
+sum()