New models, features, and fixes (#3)

- Added support for the MIT Places365 GoogleNet model.
- Added support for the Inception5h model.
- New download_models.py script for downloading the Places 365 and Inception5h models.
- Various fixes.
- New parameter for where to save CSV files when training a model.

README.md

@@ -143,6 +143,7 @@ python train_googlenet.py -data_path <training_data> -balance_classes -batch_siz
 * `-balance_classes`: Enabling this flag will balance training for each class based on class size.
 
 **Model options:**
+* `-model_file`: Path to the `.pth` model file to use for the starting model. Default is the BVLC GoogleNet model.
 * `-freeze_to`: Which layer to freeze the model up to; one of `none`, `conv1`, `conv2`, `conv3`, `mixed3a`, `mixed3b`, `mixed4a`, `mixed4b`, `mixed4c`, `mixed4d`, `mixed4e`, `mixed5a`, or `mixed5b`. Default is `mixed3b`.
 * `-freeze_aux1_to`: Which layer to freeze the first auxiliary branch up to; one of `none`, `loss_conv`, `loss_fc`, or `loss_classifier`. Default is `none`.
 * `-freeze_aux2_to`: Which layer to freeze the second auxiliary branch up to; one of `none`, `loss_conv`, `loss_fc`, or `loss_classifier`. Default is `none`.
@@ -153,6 +154,7 @@ python train_googlenet.py -data_path <training_data> -balance_classes -batch_siz
 * `-output_name`: Name of the output model. Default is `bvlc_out.pth`.
 * `-individual_acc`: Enabling this flag will print the individual accuracy of each class.
 * `-save_csv`: Enabling this flag will save loss and accuracy data to txt files.
+* `-csv_dir`: Where to save csv files. Default is set to current working directory.
 
 **Other options:**
 * `-use_device`: Zero-indexed ID of the GPU to use plus `cuda:`. Default is `cuda:0`.

models/download_models.py

@@ -0,0 +1,44 @@
+from os import path
+from sys import version_info
+if version_info[0] < 3:
+    import urllib
+else:
+    import urllib.request
+import argparse
+
+
+
+def main():
+    params = params_list()
+    if params.models == 'all' or params.models == 'places365':
+        print("Downloading the Places365 model")
+        fileurl = "https://github.com/ProGamerGov/pytorch-places/raw/restructured-models/pt_places365.pth"
+        name = "pt_places365.pth"
+        download_file(fileurl, name, params.download_path)
+    if params.models == 'all' or params.models == 'inception5h':
+        print("Downloading the Inception5h model")
+        fileurl = "https://github.com/ProGamerGov/pytorch-old-tensorflow-models/raw/restructured-models/pt_inception5h.pth"
+        name = "pt_inception5h.pth"
+        download_file(fileurl, name, params.download_path)
+
+    print("All selected models have been successfully downloaded")
+
+
+def download_file(fileurl, name, download_path):
+    if version_info[0] < 3:
+        urllib.URLopener().retrieve(fileurl, path.join(download_path, name))
+    else:
+        urllib.request.urlretrieve(fileurl, path.join(download_path, name))
+
+
+def params_list():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-models", choices=['all', 'places365', 'inception5h'], default='all')
+    parser.add_argument("-download_path", help="Download location for models", default='models')
+    params = parser.parse_args()
+    return params
+
+
+
+if __name__ == "__main__":
+    main()

train_googlenet.py

@@ -14,6 +14,7 @@ def main():
     parser.add_argument("-model_file", type=str, default='models/pt_bvlc.pth')
     parser.add_argument("-data_mean", type=str, default='')
     parser.add_argument("-data_sd", type=str, default='')
+    parser.add_argument("-base_model", choices=['bvlc', 'p365', '5h'], default='bvlc')
 
     # Training options
     parser.add_argument("-num_epochs", type=int, default=120)
@@ -29,6 +30,7 @@ def main():
     parser.add_argument("-output_name", type=str, default='bvlc_out.pth')
     parser.add_argument("-individual_acc", action='store_true')
     parser.add_argument("-save_csv", action='store_true')
+    parser.add_argument("-csv_dir", type=str, default='')
 
     # Other options
     parser.add_argument("-not_caffe", action='store_true')
@@ -69,7 +71,7 @@ def main_func(params):
 
 
     # Setup model definition
-    cnn, is_start_model = setup_model(params.model_file, num_classes=num_classes, pretrained=not params.reset_weights)
+    cnn, is_start_model, base_model = setup_model(params.model_file, num_classes=num_classes, base_model=params.base_model, pretrained=not params.reset_weights)
 
     if params.optimizer == 'sgd':
         optimizer = optim.SGD(cnn.parameters(), lr=params.lr, momentum=0.9)
@@ -145,13 +147,13 @@ def main_func(params):
         torch.backends.cudnn.enabled = True
 
 
-    save_info = [[params.data_mean, params.data_sd], num_classes, has_branches, 'bvlc_googlenet']
+    save_info = [[params.data_mean, params.data_sd], num_classes, has_branches, base_model]
 
     # Train model
     train_model(model=cnn, dataloaders=training_data, criterion=criterion, optimizer=optimizer, lrscheduler=lrscheduler, \
                 num_epochs=params.num_epochs, start_epoch=start_epoch, save_epoch=params.save_epoch, output_name=params.output_name, \
                 device=params.use_device, has_branches=has_branches, fc_only=False, num_classes=num_classes, individual_acc=params.individual_acc, \
-                should_save_csv=params.save_csv, save_info=save_info)
+                should_save_csv=params.save_csv, csv_path=params.csv_dir, save_info=save_info)
 
 
 

utils/inceptionv1_caffe.py

@@ -72,10 +72,10 @@ def __init__(self, out_features=1000, load_branches=True, mode='bvlc'):
 
         if self.mode == 'p365' or self.mode == 'bvlc':
             lrn_vals = (5, 9.999999747378752e-05, 0.75, 1)
-            diff_channels = [208, 512, 32]
+            diff_channels = [208, 512, 32, 512]
         elif self.mode == '5h':
             lrn_vals = (9, 9.99999974738e-05, 0.5, 1)
-            diff_channels = [204, 508, 48]
+            diff_channels = [204, 508, 48, 508]
 
         self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=(7, 7), stride=(2, 2), padding=(3,3), groups=1, bias=True)
         self.conv1_relu = ReluLayer()
@@ -107,7 +107,7 @@ def __init__(self, out_features=1000, load_branches=True, mode='bvlc'):
             self.fc = nn.Linear(1024, out_features)
 
         if load_branches:
-            self.aux1 = AuxBranch(512, out_features)
+            self.aux1 = AuxBranch(diff_channels[3], out_features)
             self.aux2 = AuxBranch(528, out_features)
 
 
@@ -135,6 +135,8 @@ def replace_fc(self, out_features, use_branches):
         self.fc = nn.Linear(1024, out_features)
         self.fc.apply(self.initialize_layers)
         if use_branches or self.use_branches:
+            if use_branches and self.use_branches == False:
+                self.add_branches(out_features)
             self.use_branches = True
             self.aux1.loss_classifier = nn.Linear(in_features = 1024, out_features = out_features, bias = True)
             self.aux2.loss_classifier = nn.Linear(in_features = 1024, out_features = out_features, bias = True)

utils/train_model.py

@@ -6,7 +6,7 @@
 # Training function adapted from: https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
 def train_model(model, dataloaders, criterion, optimizer=None, lrscheduler=None, num_epochs=120, start_epoch=1, save_epoch=1, \
                 output_name='out.pth', device='cuda:0', has_branches=False, fc_only=False, num_classes=0, individual_acc = False, \
-                should_save_csv=False, save_info=None):
+                should_save_csv=False, csv_path='', save_info=None):
 
     if individual_acc:
         data_counts = {}
@@ -76,7 +76,7 @@ def train_model(model, dataloaders, criterion, optimizer=None, lrscheduler=None,
             if individual_acc:
                 print('  Class Acc:', {i:j for i,j in enumerate(class_acc)})
                 if should_save_csv:
-                    save_csv_data(phase + '_acc.txt', [str(epoch)] + class_acc)
+                    save_csv_data(csv_path, phase + '_acc.txt', [str(epoch)] + class_acc)
 
             time_elapsed = time.time() - start_time
             print('  Time Elapsed {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))

utils/training_utils.py

@@ -144,28 +144,38 @@ def get_fc_channel_classes(data_path):
 
 
 # Save list as CSV
-def save_csv_data(filename, new_line):
-    with open(filename, "a+", newline='') as f:
+def save_csv_data(csv_path, filename, new_line):
+    with open(os.path.join(csv_path, filename), "a+", newline='') as f:
         wr = csv_writer(f, quoting=QUOTE_ALL)
         wr.writerow(new_line)
 
 
 # Load model class and optionally reset weights
-def setup_model(model_file='bn', num_classes=120, pretrained=False):
-    base_list = {'pt_bvlc.pth': 1000}
-    base_name = os.path.basename(model_file)
+def setup_model(model_file='pt_bvlc.pth', num_classes=120, base_model='bvlc', pretrained=False):
+    base_list = {'pt_bvlc.pth': (1000, 'bvlc'), 'pt_places365.pth': (365, 'p365'), 'pt_inception5h.pth': (1008, '5h')}
+    base_name, has_branches = os.path.basename(model_file), True
     if base_name.lower() in base_list:
-        load_classes = base_list[base_name.lower()]
+        load_classes, mode = base_list[base_name.lower()]
         is_start_model = True
+        if mode == '5h':
+            has_branches = False
     else:
         load_classes = num_classes
         is_start_model = False
+        try:
+            mode = torch.load(model_file, map_location='cpu')['base_model']
+        except:
+            mode = base_model
+        try:
+            has_branches = torch.load(model_file, map_location='cpu')['has_branches']
+        except:
+            pass
 
-    cnn = InceptionV1_Caffe(load_classes, mode='bvlc')
+    cnn = InceptionV1_Caffe(load_classes, mode=mode, load_branches=has_branches)
 
     if not pretrained:
         cnn.apply(reset_weights)
-    return cnn, is_start_model
+    return cnn, is_start_model, mode
 
 
 # Load checkpoint

utils/vis_utils.py

@@ -47,25 +47,33 @@ def simple_deprocess(output_tensor, output_name, input_mean, not_caffe):
 
 
 # Build the model definition
-def build_model(model_file='pt_places365.pth', num_classes=120, load_branches=True):
-    base_list = {'pt_places365.pth': 365, 'pt_inception5h.pth': 1008, 'pt_bvlc.pth': 1000}
+def build_model(model_file='pt_bvlc.pth', mode='bvlc', num_classes=120, load_branches=True):
+    base_list = {'pt_bvlc.pth': (1000, 'bvlc'), 'pt_places365.pth': (365, 'p365'), 'pt_inception5h.pth': (1008, '5h')}
     base_name = os.path.basename(model_file)
     if base_name.lower() in base_list:
-        load_classes = base_list[base_name.lower()]
+        load_classes, mode = base_list[base_name.lower()]
+        if mode == '5h':
+            load_branches = False
     else:
         load_classes = num_classes
 
-    cnn = InceptionV1_Caffe(load_classes, mode='bvlc', load_branches=load_branches)
+    cnn = InceptionV1_Caffe(load_classes, mode=mode, load_branches=load_branches)
 
     if base_name.lower() not in base_list:
         cnn.replace_fc(load_classes, load_branches)
     return cnn
 
 
-def load_model(model_file, num_classes=120, has_branches=True):
+def load_model(model_file, num_classes=120, has_branches=True, mode='bvlc'):
     checkpoint = torch.load(model_file, map_location='cpu')
 
     # Attempt to load model class info
+    try:
+        base_model = checkpoint['base_model']
+    except:
+        base_model = None
+    mode = mode if base_model == None else base_model
+
     try:
         loaded_classes = checkpoint['num_classes']
     except:
@@ -82,7 +90,7 @@ def load_model(model_file, num_classes=120, has_branches=True):
         h_branches = None
     has_branches = has_branches if h_branches == None else h_branches
 
-    cnn = build_model(model_file, num_classes, load_branches=has_branches)
+    cnn = build_model(model_file, mode, num_classes, load_branches=has_branches)
 
     if type(checkpoint) == dict:
         model_keys = checkpoint.keys()

vis.py

@@ -50,7 +50,7 @@ def main_func(params):
         torch.backends.cudnn.benchmark = True
         torch.backends.cudnn.enabled = True
 
-    cnn, norm_vals, num_classes = load_model(params.model_file, params.num_classes, has_branches=params.no_branches)
+    cnn, norm_vals, num_classes = load_model(params.model_file, params.num_classes, has_branches=not params.no_branches)
     if norm_vals != None:
         params.data_mean = norm_vals[0]
     else:

vis_fc.py

@@ -56,7 +56,7 @@ def main_func(params):
     except:
         model_epoch = params.model_epoch
 
-    cnn, norm_vals, num_classes = load_model(params.model_file, params.num_classes, has_branches=params.no_branches)
+    cnn, norm_vals, num_classes = load_model(params.model_file, params.num_classes, has_branches=not params.no_branches)
     if norm_vals != None:
         params.data_mean = norm_vals[0]
     else: