| 42 | and callable(torchvision_models.__dict__[name])) |
| 43 | |
| 44 | def get_args_parser(): |
| 45 | parser = argparse.ArgumentParser('DINO', add_help=False) |
| 46 | |
| 47 | # Model parameters |
| 48 | parser.add_argument('--arch', default='deit_tiny', type=str, |
| 49 | choices=['vit_tiny', 'vit_small', 'vit_base', 'deit_tiny', 'deit_small'] + torchvision_archs, |
| 50 | help="""Name of architecture to train. For quick experiments with ViTs, |
| 51 | we recommend using vit_tiny or vit_small.""") |
| 52 | parser.add_argument('--patch_size', default=16, type=int, help="""Size in pixels |
| 53 | of input square patches - default 16 (for 16x16 patches). Using smaller |
| 54 | values leads to better performance but requires more memory. Applies only |
| 55 | for ViTs (vit_tiny, vit_small and vit_base). If <16, we recommend disabling |
| 56 | mixed precision training (--use_fp16 false) to avoid unstabilities.""") |
| 57 | parser.add_argument('--out_dim', default=65536, type=int, help="""Dimensionality of |
| 58 | the DINO head output. For complex and large datasets large values (like 65k) work well.""") |
| 59 | parser.add_argument('--out_dim_selfpatch', default=4096, type=int, help="""Dimensionality of |
| 60 | the DINO head output. For complex and large datasets large values (like 65k) work well.""") |
| 61 | parser.add_argument('--norm_last_layer', default=True, type=utils.bool_flag, |
| 62 | help="""Whether or not to weight normalize the last layer of the DINO head. |
| 63 | Not normalizing leads to better performance but can make the training unstable. |
| 64 | In our experiments, we typically set this paramater to False with vit_small and True with vit_base.""") |
| 65 | parser.add_argument('--momentum_teacher', default=0.996, type=float, help="""Base EMA |
| 66 | parameter for teacher update. The value is increased to 1 during training with cosine schedule. |
| 67 | We recommend setting a higher value with small batches: for example use 0.9995 with batch size of 256.""") |
| 68 | parser.add_argument('--use_bn_in_head', default=False, type=utils.bool_flag, |
| 69 | help="Whether to use batch normalizations in projection head (Default: False)") |
| 70 | parser.add_argument("--k_num", default=4, type=int, help="top k confident patch") |
| 71 | |
| 72 | # Temperature teacher parameters |
| 73 | parser.add_argument('--warmup_teacher_temp', default=0.04, type=float, |
| 74 | help="""Initial value for the teacher temperature: 0.04 works well in most cases. |
| 75 | Try decreasing it if the training loss does not decrease.""") |
| 76 | parser.add_argument('--teacher_temp', default=0.04, type=float, help="""Final value (after linear warmup) |
| 77 | of the teacher temperature. For most experiments, anything above 0.07 is unstable. We recommend |
| 78 | starting with the default value of 0.04 and increase this slightly if needed.""") |
| 79 | parser.add_argument('--warmup_teacher_temp_epochs', default=0, type=int, |
| 80 | help='Number of warmup epochs for the teacher temperature (Default: 30).') |
| 81 | |
| 82 | # Training/Optimization parameters |
| 83 | parser.add_argument('--use_fp16', type=utils.bool_flag, default=True, help="""Whether or not |
| 84 | to use half precision for training. Improves training time and memory requirements, |
| 85 | but can provoke instability and slight decay of performance. We recommend disabling |
| 86 | mixed precision if the loss is unstable, if reducing the patch size or if training with bigger ViTs.""") |
| 87 | parser.add_argument('--weight_decay', type=float, default=0.04, help="""Initial value of the |
| 88 | weight decay. With ViT, a smaller value at the beginning of training works well.""") |
| 89 | parser.add_argument('--weight_decay_end', type=float, default=0.4, help="""Final value of the |
| 90 | weight decay. We use a cosine schedule for WD and using a larger decay by |
| 91 | the end of training improves performance for ViTs.""") |
| 92 | parser.add_argument('--clip_grad', type=float, default=3.0, help="""Maximal parameter |
| 93 | gradient norm if using gradient clipping. Clipping with norm .3 ~ 1.0 can |
| 94 | help optimization for larger ViT architectures. 0 for disabling.""") |
| 95 | parser.add_argument('--batch_size_per_gpu', default=64, type=int, |
| 96 | help='Per-GPU batch-size : number of distinct images loaded on one GPU.') |
| 97 | parser.add_argument('--epochs', default=300, type=int, help='Number of epochs of training.') |
| 98 | parser.add_argument('--freeze_last_layer', default=1, type=int, help="""Number of epochs |
| 99 | during which we keep the output layer fixed. Typically doing so during |
| 100 | the first epoch helps training. Try increasing this value if the loss does not decrease.""") |
| 101 | parser.add_argument("--lr", default=0.0005, type=float, help="""Learning rate at the end of |