feat: Add NaViT

.gitignore

@@ -172,3 +172,4 @@ rank_*/
 ckpt/
 output/
 outputs/
+kernel_meta*/

config.py

@@ -126,6 +126,9 @@ def create_parser():
                        help='Number of dataset repetition for repeated augmentation. '
                             'If 0 or 1, repeated augmentation is disabled. '
                             'Otherwise, repeated augmentation is enabled and the common choice is 3 (default=0)')
+    group.add_argument('--patch_size', type=int, default=32, help="Patch size in sequence packing.")
+    group.add_argument('--max_seq_length', type=int, default=2048, help="maximum sequence length in sequence packing.")
+    group.add_argument('--max_num_each_group', type=int, default=40, help="maximum number of images in each sequence")
 
     # Model parameters
     group = parser.add_argument_group('Model parameters')

configs/navit/navit_b16_384_ascend.yaml

@@ -0,0 +1,55 @@
+# system
+mode: 0
+distribute: True
+num_parallel_workers: 8
+val_while_train: True
+val_interval: 1
+val_split: val
+
+# dataset
+dataset: "imagenet"
+data_dir: "/path/to/imagenet"
+shuffle: True
+dataset_download: False
+batch_size: 64
+drop_remainder: True
+patch_size: 16
+max_seq_length: 768
+max_num_each_group: 16
+
+# augmentation
+image_resize: 384
+hflip: 0.5
+interpolation: "bicubic"
+
+# model
+model: "navit_b_16_384"
+drop_rate: 0.1
+drop_path_rate: 0.1
+num_classes: 1000
+pretrained: False
+ckpt_path: ""
+keep_checkpoint_max: 1
+ckpt_save_policy: "top_k"
+ckpt_save_dir: "./ckpt"
+epoch_size: 100
+dataset_sink_mode: False
+amp_level: "O2"
+
+# loss
+loss: "TOKEN_CE"
+loss_scale: 1024.0
+label_smoothing: 0.1
+
+# lr scheduler
+scheduler: "warmup_cosine_decay"
+lr: 0.0001
+min_lr: 1e-6
+warmup_epochs: 10
+decay_epochs: 90
+lr_epoch_stair: False
+
+# optimizer
+opt: "adamw"
+weight_decay: 0.1
+use_nesterov: False