pytorch实现moco v3

链接：MoCo_v3_pytorch

参考论文：An Empirical Study of Training Self-Supervised Vision Transformers

训练

CUDA_VISIBLE_DEVICES=6 python main_mocov3.py -a resnet18 --lr 0.03 --batch-size 1 --dist-url 'tcp://localhost:10003' --multiprocessing-distributed --world-size 1 --rank 0 --moco-k 65536  --epochs 10 /your_data_root

数据集

例子中采用了一个toy数据集，随机产生一张256x256大小的图片，用户可以根据需要替换成自己需要的数据集。

只需要修改main_mocov3.py文件中main_worker中的train_dataset = ToyDS()行，替换成自己的数据集即可。

VIT模块参数修改

修改main_mocov3.py文件中main_worker中的

    v = ViT(
        image_size = 256,
        patch_size = 32,
        num_classes = 10,
        dim = 1024,
        depth = 6,
        heads = 16,
        mlp_dim = 2048,
        dropout = 0.1,
        emb_dropout = 0.1
    )

具体参数含义可以参考An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale