From 50bfbee7a0f7fb54178538fb2e913ea468e47378 Mon Sep 17 00:00:00 2001 From: yjh0410 <1394571815@qq.com> Date: Tue, 22 Aug 2023 13:30:55 +0800 Subject: [PATCH] train MAE-ViT-Nano on CIFAR --- README.md | 12 ++++++------ mae_pretrain.py | 6 +++--- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 6b4b671..00678b5 100644 --- a/README.md +++ b/README.md @@ -6,39 +6,39 @@ PyTorch implementation of Masked AutoEncoder - Train `MAE-ViT-Nano` on CIFAR10 dataset: ```Shell -python mae_pretrain.py --dataset cifar10 -m mae_vit_nano --batch_size 256 --img_size 32 --patch_size 2 +python mae_pretrain.py --dataset cifar10 -m mae_vit_nano --batch_size 256 --img_size 32 --patch_size 2 --max_epoch 400 --wp_epoch 40 ``` - Train `MAE-ViT-Nano` on ImageNet dataset: ```Shell -python mae_finetune.py --dataset imagenet -m mae_vit_nano --batch_size 256 --img_size 224 --patch_size 16 +python mae_finetune.py --dataset imagenet -m mae_vit_nano --batch_size 256 --img_size 224 --patch_size 16 --max_epoch 400 --wp_epoch 40 ``` ## 2. Train from scratch - Train `ViT-Nano` on CIFAR10 dataset: ```Shell -python mae_finetune.py --dataset cifar10 -m vit_nano --batch_size 256 --img_size 32 --patch_size 2 +python mae_finetune.py --dataset cifar10 -m vit_nano --batch_size 256 --img_size 32 --patch_size 2 --max_epoch 200 --wp_epoch 20 ``` - Train `ViT-Nano` on ImageNet dataset: ```Shell -python mae_finetune.py --dataset imagenet -m vit_nano --batch_size 256 --img_size 224 --patch_size 16 +python mae_finetune.py --dataset imagenet -m vit_nano --batch_size 256 --img_size 224 --patch_size 16 --max_epoch 200 --wp_epoch 20 ``` ## 3. Train from MAE pretrained - Train `ViT-Nano` on CIFAR10 dataset: ```Shell -python mae_finetune.py --dataset cifar10 -m vit_nano --batch_size 256 --img_size 32 --patch_size 2 --mae_pretrained +python mae_finetune.py --dataset cifar10 -m vit_nano --batch_size 256 --img_size 32 --patch_size 2 --mae_pretrained --max_epoch 50 --wp_epoch 5 ``` - Train `ViT-Nano` on ImageNet dataset: ```Shell -python mae_finetune.py --dataset imagenet -m vit_nano --batch_size 256 --img_size 224 --patch_size 16 --mae_pretrained +python mae_finetune.py --dataset imagenet -m vit_nano --batch_size 256 --img_size 224 --patch_size 16 --mae_pretrained --max_epoch 50 --wp_epoch 5 ``` ## 4. Experiments diff --git a/mae_pretrain.py b/mae_pretrain.py index 94ba2a4..3f80ea8 100644 --- a/mae_pretrain.py +++ b/mae_pretrain.py @@ -201,7 +201,7 @@ def main(): loss = mae_loss(images, output['x_pred'], output['mask'], args.patch_size, args.norm_pix_loss) # update num_fgs & losses total_num_fgs += output['mask'].sum().item() - total_losses += loss * output['mask'].sum().item() + total_losses += loss.item() * output['mask'].sum().item() # Backward loss /= args.grad_accumulate @@ -255,8 +255,8 @@ def main(): 'optimizer': optimizer.state_dict(), 'epoch': epoch}, checkpoint_path) - total_num_fgs = 0. - total_losses = 0. + total_num_fgs = 0. + total_losses = 0. lr_scheduler.step()