works for cifar

NVlabs · May 2, 2022 · 941e924 · 941e924
1 parent 68c0000
commit 941e924
Show file tree

Hide file tree

Showing 13 changed files with 5 additions and 3 deletions.
diff --git a/classification/__pycache__/afnonet.cpython-36.pyc b/classification/__pycache__/afnonet.cpython-36.pyc
diff --git a/classification/__pycache__/main_afnonet.cpython-36.pyc b/classification/__pycache__/main_afnonet.cpython-36.pyc
diff --git a/classification/afnonet.py b/classification/afnonet.py
@@ -140,7 +140,8 @@ def __init__(self, dim, mlp_ratio=4., drop=0., drop_path=0., act_layer=nn.GELU,
 
         #to be added soon ... @John: pls double check
         if args.mixing_type == "afno":
-            self.filter = AFNO2D(hidden_size=768, num_blocks=8, sparsity_threshold=0.01, hard_thresholding_fraction=1, hidden_size_factor=1)
+            self.filter = AFNO2D(hidden_size=args.hidden_size, num_blocks=args.fno_blocks, sparsity_threshold=0.01, hard_thresholding_fraction=1, hidden_size_factor=1)
+            #self.filter = AdaptiveFourierNeuralOperator(dim, h=h, w=w)
         elif args.mixing_type == "bfno":
             self.filter = BFNO2D(hidden_size=768, num_blocks=8, hard_thresholding_fraction=1)
         elif args.mixing_type == "sa":

diff --git a/classification/checkpoint_best.pth b/classification/checkpoint_best.pth
diff --git a/classification/main_afnonet.py b/classification/main_afnonet.py
@@ -142,7 +142,7 @@ def get_args_parser():
     # Dataset parameters
     parser.add_argument('--data-path', default='/datasets01/imagenet_full_size/061417/', type=str,
                         help='dataset path')
-    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR', 'IMNET', 'INAT', 'INAT19'],
+    parser.add_argument('--data-set', default='IMNET', choices=['CIFAR10', 'CIFAR100', 'CARS', 'FLOWERS', 'IMNET'],
                         type=str, help='Image Net dataset path')
     parser.add_argument('--inat-category', default='name',
                         choices=['kingdom', 'phylum', 'class', 'order', 'supercategory', 'family', 'genus', 'name'],
@@ -229,6 +229,7 @@ def main(args):
     cudnn.benchmark = True
 
     dataset_train, args.nb_classes = build_dataset(is_train=True, args=args)
+    print('args.nb_classes', args.nb_classes)
     dataset_val, _ = build_dataset(is_train=False, args=args)
 
     if True:  # args.distributed:
@@ -283,7 +284,7 @@ def main(args):
         print('mix up is not used')
 
     model = AFNONet(
-        img_size=args.input_size, 
+        img_size=args.input_size, num_classes=args.nb_classes,
         patch_size=4, embed_dim=args.hidden_size, depth=args.num_layers, mlp_ratio=4,
         norm_layer=partial(nn.LayerNorm, eps=1e-6)
     )

diff --git a/...1-26_rno1-m03-g01-dgx1-003/events.out.tfevents.1651447286.rno1-m03-g01-dgx1-003.2475198.0 b/...1-26_rno1-m03-g01-dgx1-003/events.out.tfevents.1651447286.rno1-m03-g01-dgx1-003.2475198.0
diff --git a/...0-19_rno1-m03-g01-dgx1-003/events.out.tfevents.1651447819.rno1-m03-g01-dgx1-003.2478944.0 b/...0-19_rno1-m03-g01-dgx1-003/events.out.tfevents.1651447819.rno1-m03-g01-dgx1-003.2478944.0
diff --git a/...4-02_rno1-m03-g01-dgx1-003/events.out.tfevents.1651448642.rno1-m03-g01-dgx1-003.2492330.0 b/...4-02_rno1-m03-g01-dgx1-003/events.out.tfevents.1651448642.rno1-m03-g01-dgx1-003.2492330.0
diff --git a/...9-10_rno1-m03-g01-dgx1-003/events.out.tfevents.1651448950.rno1-m03-g01-dgx1-003.2494606.0 b/...9-10_rno1-m03-g01-dgx1-003/events.out.tfevents.1651448950.rno1-m03-g01-dgx1-003.2494606.0
diff --git a/...7-15_rno1-m03-h08-dgx1-095/events.out.tfevents.1651473435.rno1-m03-h08-dgx1-095.2777745.0 b/...7-15_rno1-m03-h08-dgx1-095/events.out.tfevents.1651473435.rno1-m03-h08-dgx1-095.2777745.0
diff --git a/...1-50_rno1-m03-h08-dgx1-095/events.out.tfevents.1651473710.rno1-m03-h08-dgx1-095.2779573.0 b/...1-50_rno1-m03-h08-dgx1-095/events.out.tfevents.1651473710.rno1-m03-h08-dgx1-095.2779573.0
diff --git a/...3-21_rno1-m03-h08-dgx1-095/events.out.tfevents.1651473801.rno1-m03-h08-dgx1-095.2780765.0 b/...3-21_rno1-m03-h08-dgx1-095/events.out.tfevents.1651473801.rno1-m03-h08-dgx1-095.2780765.0
diff --git a/...6-51_rno1-m03-h08-dgx1-095/events.out.tfevents.1651474611.rno1-m03-h08-dgx1-095.2787752.0 b/...6-51_rno1-m03-h08-dgx1-095/events.out.tfevents.1651474611.rno1-m03-h08-dgx1-095.2787752.0