bug fix in issue #57

ahmetgunduz · May 13, 2020 · f5ded06 · f5ded06
1 parent 980545b
commit f5ded06
Showing 1 changed file with 49 additions and 6 deletions.
diff --git a/model.py b/model.py
@@ -159,12 +159,27 @@ def generate_model(opt):
             else:
                 model.module.fc = nn.Linear(model.module.fc.in_features, opt.n_finetune_classes)
                 model.module.fc = model.module.fc.cuda()
+        if opt.modality == 'RGB' and opt.model != 'c3d':
+            print("[INFO]: RGB model is used for init model")
+            model = _modify_first_conv_layer(model,3,3) ##### Check models trained (3,7,7) or (7,7,7)
+        elif opt.modality == 'Depth':
+            print("[INFO]: Converting the pretrained model to Depth init model")
+            model = _construct_depth_model(model)
+            print("[INFO]: Done. Flow model ready.")
+        elif opt.modality == 'RGB-D':
+            print("[INFO]: Converting the pretrained model to RGB+D init model")
+            model = _construct_rgbdepth_model(model)
+            print("[INFO]: Done. RGB-D model ready.")
 
-            # model = _modify_first_conv_layer(model)
-            # model = model.cuda()
+        modules = list(model.modules())
+        first_conv_idx = list(filter(lambda x: isinstance(modules[x], nn.Conv3d),
+                                                   list(range(len(modules)))))[0]
+        conv_layer = modules[first_conv_idx]
+        if conv_layer.kernel_size[0]> opt.sample_duration:
+            model = _modify_first_conv_layer(model,int(opt.sample_duration/2),1) 
 
-            parameters = get_fine_tuning_parameters(model, opt.ft_portion)
-            return model, parameters
+        parameters = get_fine_tuning_parameters(model, opt.ft_portion)
+        return model, parameters
     else:
         if opt.pretrain_path:
             print('loading pretrained model {}'.format(opt.pretrain_path))
@@ -186,8 +201,36 @@ def generate_model(opt):
             else:
                 model.module.fc = nn.Linear(model.module.fc.in_features, opt.n_finetune_classes)
 
-            parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
-            return model, parameters
+        if opt.modality == 'RGB' and opt.model != 'c3d':
+            print("[INFO]: RGB model is used for init model")
+            model = _modify_first_conv_layer(model,3,3)
+        elif opt.modality == 'Depth':
+            print("[INFO]: Converting the pretrained model to Depth init model")
+            model = _construct_depth_model(model)
+            print("[INFO]: Deoth model ready.")
+        elif opt.modality == 'RGB-D':
+            print("[INFO]: Converting the pretrained model to RGB-D init model")
+            model = _construct_rgbdepth_model(model)
+            print("[INFO]: Done. RGB-D model ready.")
+
+        modules = list(model.modules())
+        first_conv_idx = list(filter(lambda x: isinstance(modules[x], nn.Conv3d),
+                                                   list(range(len(modules)))))[0]
+        conv_layer = modules[first_conv_idx]
+        if conv_layer.kernel_size[0]> opt.sample_duration:
+            print("[INFO]: RGB model is used for init model")
+            model = _modify_first_conv_layer(model,int(opt.sample_duration/2),1) 
+
+        if opt.model == 'c3d':# CHECK HERE
+            model.fc = nn.Linear(
+                model.fc[0].in_features, model.fc[0].out_features)
+        else:
+            model.fc = nn.Linear(model.fc.in_features,
+                                        opt.n_finetune_classes)
+
+        parameters = get_fine_tuning_parameters(model, opt.ft_begin_index)
+        return model, parameters
+
 
     return model, model.parameters()