MaJ example simple + Suppression Nesterov momentum

higher/smart_aug/smart_aug_example.py

@@ -2,31 +2,12 @@
 
 """
 
-from model import *
+from LeNet import *
 from dataug import *
 from train_utils import *
 
-# Use available TF (see transformations.py)
-tf_names = [
-    ## Geometric TF ##
-    'Identity',
-    'FlipUD',
-    'FlipLR',
-    'Rotate',
-    'TranslateX',
-    'TranslateY',
-    'ShearX',
-    'ShearY',
-
-    ## Color TF (Expect image in the range of [0, 1]) ##
-    'Contrast',
-    'Color',
-    'Brightness',
-    'Sharpness',
-    'Posterize',
-    'Solarize', #=>Image entre [0,1] #Pas opti pour des batch
-]
-
+tf_config='../config/base_tf_config.json'
+TF_loader=TF_loader()
 
 device = torch.device('cuda') #Select device to use
 
@@ -48,19 +29,19 @@ if __name__ == "__main__":
         },
         'Inner':{
             'optim': 'SGD',
-            'lr':1e-2, #1e-2
+            'lr':1e-2, #1e-2/1e-1 (ResNet)
             'momentum':0.9, #0.9
+            'decay':0.0005, #0.0005
+            'nesterov':False, #False (True: Bad behavior w/ Data_aug)
+            'scheduler':'cosine', #None, 'cosine', 'multiStep', 'exponential'
         }
     }
 
     #Models
     model = LeNet(3,10)
-    #model = ResNet(num_classes=10)
-    #model = MobileNetV2(num_classes=10)
-    #model = WideResNet(num_classes=10, wrn_size=32)
 
     #Smart_aug initialisation
-    tf_dict = {k: TF.TF_dict[k] for k in tf_names}
+    tf_dict, tf_ignore_mag =TF_loader.load_TF_dict(tf_config)
     model = Higher_model(model) #run_dist_dataugV3
     aug_model = Augmented_model(
         Data_augV5(TF_dict=tf_dict, 
@@ -68,7 +49,8 @@ if __name__ == "__main__":
             mix_dist=0.8, 
             fixed_prob=False, 
             fixed_mag=False, 
-            shared_mag=False), 
+            shared_mag=False,
+            TF_ignore_mag=tf_ignore_mag), 
         model).to(device)
 
     print("{} on {} for {} epochs - {} inner_it".format(str(aug_model), device_name, epochs, n_inner_iter))