cnn config update

Author: UltraShieldRog

code/algorithm/main.py

@@ -127,8 +127,12 @@ def main(config):
             model = TFIDF(ImageCaptionDataset.CLASSES,
                           tfidf_vectorizer,
                           config.threshold)
-        elif config.model_type == 'cnn':
-            model = CNN(ImageCaptionDataset.CLASSES, 'alexnet')
+        elif config.model_type in [
+                'resnet', 'alexnet', 'vgg',
+                'densenet', 'googlenet', 'resnext',
+                'wide_resnet', 'mnasnet'
+                ]:
+            model = CNN(ImageCaptionDataset.CLASSES, config.model_type)
 
         loss_func = BCEWithLogitsLoss()
         optimiser = Adam(model.parameters())
@@ -433,7 +437,14 @@ def parse_args(args):
                  'rcnn',
                  'lstm',
                  'tfidf',
-                 'cnn'
+                 'resnet',
+                 'alexnet',
+                 'vgg',
+                 'densenet',
+                 'googlenet',
+                 'resnext',
+                 'wide_resnet',
+                 'mnasnet'
                  ],
         type=str,
         required=False,

code/logs/20200529-021503.txt

@@ -0,0 +1,2 @@
+micro_f1: 0.702053, macro_f1: 0.491366, weighted_f1: 0.710688
+micro_f1: 0.710000, macro_f1: 0.475064, weighted_f1: 0.727771