MCTformer: Multi-class Token Transformer

The main.py script serves as the primary entrypoint for training MCTformer models or performing evaluation and attention map generation. It uses argparse to handle a wide range of hyperparameters for model architecture, optimization, augmentation, and dataset configuration.

To use the script, you must provide arguments via the command line. The script supports three primary modes:

1. Training: The default mode. It runs for a specified number of --epochs and saves checkpoints to --output_dir.
2. Evaluation: Triggered by the --eval flag. It performs evaluation on the validation dataset using the model loaded from --resume.
3. Attention Map Generation: Triggered by the --gen_attention_maps flag. It uses a checkpoint from --resume to generate and save attention maps.

Note: When finetuning, the script includes logic to automatically interpolate position embeddings and handle class token mismatches if the pretrained checkpoint has a different number of classes or input resolution.

# Example: Training a model
python main.py --model deit_small_MCTformerPlus --batch-size 32 --epochs 45 --data-path /path/to/VOC2012 --output_dir ./saved_model

# Example: Evaluating a trained model
python main.py --eval --resume ./saved_model/checkpoint.pth --data-path /path/to/VOC2012

# Example: Generating attention maps
python main.py --gen_attention_maps --resume ./saved_model/checkpoint.pth --data-path /path/to/VOC2012

The following arguments control the model architecture and the optimization process in main.py.

# Model parameters
--model MODEL          # Name of model to train (e.g., 'deit_small_MCTformerPlus')
--input-size INT       # Images input size
--drop PCT             # Dropout rate
--drop-path PCT       # Drop path rate

# Optimizer parameters
--opt OPTIMIZER        # Optimizer (default: 'adamw')
--opt-eps EPSILON     # Optimizer Epsilon
--opt-betas BETA      # Optimizer Betas (nargs='+')
--clip-grad NORM       # Clip gradient norm
--momentum M           # SGD momentum
--weight-decay FLOAT   # Weight decay

# Learning rate schedule parameters
--sched SCHEDULER      # LR scheduler (default: 'cosine')
--lr LR                # Learning rate
--lr-noise PCT PCT     # Learning rate noise on/off epoch percentages (nargs='+')
--lr-noise-pct PERCENT # Learning rate noise limit percent
--lr-noise-std STDDEV  # Learning rate noise std-dev
--warmup-lr LR         # Warmup learning rate
--min-lr LR            # Lower lr bound for cyclic schedulers
--decay-epochs N       # Epoch interval to decay LR
--warmup-epochs N      # Epochs to warmup LR
--cooldown-epochs N    # Epochs to cooldown LR
--patience-epochs N    # Patience epochs for Plateau LR scheduler
--decay-rate RATE       # LR decay rate

The following arguments control data augmentation strategies and dataset loading in main.py.

# Augmentation parameters
--color-jitter PCT     # Color jitter factor
--aa NAME              # AutoAugment policy (e.g., 'v0', 'original', or 'rand-m9-mstd0.5-inc1')
--smoothing FLOAT      # Label smoothing
--train-interpolation STR # Training interpolation (random, bilinear, bicubic)
--repeated-aug         # Enable repeated augmentation
--no-repeated-aug      # Disable repeated augmentation
--reprob PCT           # Random erase prob
--remode STR           # Random erase mode (e.g., 'pixel')
--recount INT          # Random erase count
--resplit               # Do not random erase first (clean) augmentation split

# Dataset parameters
--data-path STR         # Dataset path
--img-list STR         # Image list path
--data-set STR         # Dataset name (e.g., 'VOC12')
--finetune STR          # Finetune from checkpoint (URL or local path)
--resume STR           # Resume from checkpoint
--start_epoch N        # Start epoch
--eval                 # Perform evaluation only
--num_workers INT      # Number of workers for DataLoader
--pin-mem              # Pin CPU memory in DataLoader
--no-pin-mem           # Do not pin CPU memory