MoCo代码分析 [自监督学习]

✅

个人简介：南京邮电大学，计算机科学与技术，在读本科

● 前驱文章一：《MoCo v1 文献研究 [自监督学习]》

● 前驱文章二：《MoCo v2 文献研究 [自监督学习]》

文章目录

• 一、代码中的 “ReadMe.md”
• • 1.1 Preparation
  • 1.2 Unsupervised Training
  • 1.3 Linear Classification
  • 1.4 Models
  • 1.5 Transferring to Object Detection
  • 1.6 License
  • 1.7 See Also
• 二、main_moco.py —— 无监督训练主文件
• • 2.1 关于“导入包”的代码
  • 2.2 关于“命令行参数”的代码
  • 2.3 关于“main()函数”的代码
  • 2.4 关于“main_worker()函数”的代码
  • 2.5 关于“train()函数”的代码
  • 2.6 关于“AverageMeter类”的代码
  • 2.7 关于“ProgressMeter类”的代码
  • 2.8 关于“adjust_learning_rate()函数”的代码
  • 2.9 关于“accuracy()函数”的代码
  • 2.10 最后一点
• 三、builder.py —— MoCo模型建立的文件
• • 3.1 关于“MoCo类”的代码
  • 3.2 关于“concat_all_gather()函数”的代码
• 四、loader.py —— 辅助图片载入的文件
• • 4.1 关于“TwoCropsTransform类”的代码
  • 4.2 关于“GaussianBlur类”的代码
• 五、main_lincls.py —— 评估模型建立的文件
• ▶本文参考附录

**源码地址**: https://github.com/facebookresearch/moco

一、代码中的 “ReadMe.md”

● 这里简单把 “ReadMe.md” 复制了下来，方便后面输入命令来上机跑。

MoCo: Momentum Contrast for Unsupervised Visual Representation Learning

This is a PyTorch implementation of the MoCo paper:

@Article{he2019moco,author= {Kaiming He and Haoqi Fan and Yuxin Wu and Saining Xie and Ross Girshick},title = {Momentum Contrast for Unsupervised Visual Representation Learning},journal = {arXiv preprint arXiv:1911.05722},year= {2019},}

It also includes the implementation of the MoCo v2 paper:

@Article{chen2020mocov2,author= {Xinlei Chen and Haoqi Fan and Ross Girshick and Kaiming He},title = {Improved Baselines with Momentum Contrastive Learning},journal = {arXiv preprint arXiv:2003.04297},year= {2020},}

1.1 Preparation

Install PyTorch and ImageNet dataset following the official PyTorch ImageNet training code.

This repo aims to be minimal modifications on that code. Check the modifications by:

diff main_moco.py <(curl https://raw.githubusercontent.com/pytorch/examples/master/imagenet/main.py)diff main_lincls.py <(curl https://raw.githubusercontent.com/pytorch/examples/master/imagenet/main.py)

1.2 Unsupervised Training

This implementation only supports multi-gpu, DistributedDataParallel training, which is faster and simpler; single-gpu or DataParallel training is not supported.

To do unsupervised pre-training of a ResNet-50 model on ImageNet in an 8-gpu machine, run:

python main_moco.py \-a resnet50 \--lr 0.03 \--batch-size 256 \--dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \[your imagenet-folder with train and val folders]

This script uses all the default hyper-parameters as described in the MoCo v1 paper. To run MoCo v2, set --mlp --moco-t 0.2 --aug-plus --cos.

Note: for 4-gpu training, we recommend following the linear lr scaling recipe: --lr 0.015 --batch-size 128 with 4 gpus. We got similar results using this setting.

1.3 Linear Classification

With a pre-trained model, to train a supervised linear classifier on frozen features/weights in an 8-gpu machine, run:

python main_lincls.py \-a resnet50 \--lr 30.0 \--batch-size 256 \--pretrained [your checkpoint path]/checkpoint_0199.pth.tar \--dist-url 'tcp://localhost:10001' --multiprocessing-distributed --world-size 1 --rank 0 \[your imagenet-folder with train and val folders]

Linear classification results on ImageNet using this repo with 8 NVIDIA V100 GPUs :

Here we run 5 trials (of pre-training and linear classification) and report mean±std: the 5 results of MoCo v1 are {60.6, 60.6, 60.7, 60.9, 61.1}, and of MoCo v2 are {67.7, 67.6, 67.4, 67.6, 67.3}.

1.4 Models

Our pre-trained ResNet-50 models can be downloaded as following:

1.5 Transferring to Object Detection

See ./detection.

1.6 License

This project is under the CC-BY-NC 4.0 license. See LICENSE for details.

1.7 See Also

• moco.tensorflow: A TensorFlow re-implementation.
• Colab notebook: CIFAR demo on Colab GPU.

二、main_moco.py —— 无监督训练主文件

● 整个无监督训练模型的代码框架如下：

import ...parser.add_argument(...)def main():...def main_worker(gpu, ngpus_per_node, args):...def train(train_loader, model, criterion, optimizer, epoch, args):...def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):...class AverageMeter(object):"""Computes and stores the average and current value"""def __init__(self, name, fmt=':f'):...def reset(self):...def update(self, val, n=1):...def __str__(self):...class ProgressMeter(object):def __init__(self, num_batches, meters, prefix=""):...def display(self, batch):...def _get_batch_fmtstr(self, num_batches):...def adjust_learning_rate(optimizer, epoch, args):"""Decay the learning rate based on schedule"""...def accuracy(output, target, topk=(1,)):"""Computes the accuracy over the k top predictions for the specified values of k"""...if __name__ == '__main__':main()

● 接下来，一个一个地细剖。按着顺序来。

2.1 关于“导入包”的代码

● 作者把包的引入分成了三类：基础库、Pytorch库、个人写的moco库。

● 作者的模型是用 PyTorch 实现的。python torch 又称 PyTorach，是一个以 Python 优先的深度学习框架，一个开源的 Python 机器学习库，用于自然语言处理等应用程序，不仅能够实现强大的 GPU 加速，同时还支持动态神经网络，这是现在很多主流框架比如 Tensorflow 等都不支持的。

● moco.loader 和 moco.builder 是在 “moco” 文件夹下的两个子文件，分别是 “loader.py” 和 “builder.py”。

import argparseimport builtinsimport mathimport osimport randomimport shutilimport timeimport warningsimport torchimport torch.nn as nnimport torch.nn.parallelimport torch.backends.cudnn as cudnnimport torch.distributed as distimport torch.optimimport torch.multiprocessing as mpimport torch.utils.dataimport torch.utils.data.distributedimport torchvision.transforms as transformsimport torchvision.datasets as datasetsimport torchvision.models as modelsimport moco.loaderimport moco.builder

◆ 代码说明：“import ... as ...” 的意思是 “取别名”，比如对于 import torch.nn as nn 来说，后面你想调用 “torch.nn” 这个函数，只需要写 “nn” 即可。

2.2 关于“命令行参数”的代码

● 接下来所做的事情都是关于 “import argparse” 后做的事情。

● argparse 是 python 自带的命令行参数解析包，可以用来方便地读取命令行参数。它的使用也比较简单。

model_names = sorted(name for name in models.__dict__ if name.islower() and not name.startswith("__") and callable(models.__dict__[name]))parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')

◆ 代码说明：
 ① models 是从 “import torchvision.models as models” 这里来的。torchvision.models 中为我们提供了已经训练好的模型，让我们可以加载之后，直接使用。
 ② 类的静态函数、类函数、普通函数、全局变量以及一些内置的属性都是放在 “类的__dict__” 里的。而 “对象的__dict__” 中存储了一些 self.xxx 等一些东西。
 ③ islower()的功能：检查文本中所有字符是否都小写。
 ④ startswith()的功能：判断字符串是否以指定字符或字符串开头。
 ⑤ callable()的功能：检查一个对象是否是可调用的。
 ⑥ sorted()的功能：返回重新排序的列表。（默认是升序）
 ⑦ model_names = sored(...) 所做的事情就是把 torchvision.models 中的 “models.__dict__” 中的每一个 “大成员” 中的 “所有小成员” 进行了重排(除开 “字符串开头为__的” 小成员名)
 ⑧ 导入 argparse 这个包后，作者调用了包中的 “ArgumentParser类” 生成了一个 parser 对象（好多博客中把这个叫做参数解析器），其中的 description 描述这个参数解析器是干什么的，当我们在命令行显示帮助信息的时候会看到 description 描述的信息。最后作者将其赋值给了 parser。

● 要看懂下面三段代码，看一篇这个即可：官方的 argparse 用法

● 函数 add_argument(name or flags...[, action][, nargs][, const][, default][, type][, choices][, required][, help][, metavar][, dest]) 的功能：定义单个的命令行参数应当如何解析。每个形参都在下面有它自己更多的描述，代码中用到的有：
 ① name or flags：一个命名或者一个选项字符串的列表，例如 foo 或 -f, --foo。
 ② action：当参数在命令行中出现时使用的动作基本类型。
 ③ nargs：命令行参数应当消耗的数目。
 ④ default：当参数未在命令行中出现并且也不存在于命名空间对象时所产生的值。
 ⑤ type：命令行参数应当被转换成的类型。
 ⑥ choices：可用的参数的容器。
 ⑦ help：一个此选项作用的简单描述。
 ⑧ metavar：在使用方法消息中使用的参数值示例。
 ⑨ dest：被添加到 parse_args() 所返回对象上的属性名。

# 训练数据的目录parser.add_argument('data', metavar='DIR', help='path to dataset')# 选择训练模型parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet50', choices=model_names,help='model architecture: ' + ' | '.join(model_names) + ' (default: resnet50)')# 还不太清楚...后期补充parser.add_argument('-j', '--workers', default=32, type=int, metavar='N', help='number of data loading workers (default: 32)')# 迭代次数parser.add_argument('--epochs', default=200, type=int, metavar='N', help='number of total epochs to run')# 继承的迭代轮数(默认为0), 也就是说, 可以持续训练, 上一天关机保存的 模型, 可继续加载, 继续训练, 只不过迭代的轮数要写入(不然默认为0)parser.add_argument('--start-epoch', default=0, type=int, metavar='N', help='manual epoch number (useful on restarts)')# 批量大小parser.add_argument('-b', '--batch-size', default=256, type=int, metavar='N',help='mini-batch size (default: 256), this is the total ' 'batch size of all GPUs on the current node when ' 'using Data Parallel or Distributed Data Parallel')# 学习速率parser.add_argument('--lr', '--learning-rate', default=0.03, type=float, metavar='LR', help='initial learning rate', dest='lr')# 还不太清楚...后期补充parser.add_argument('--schedule', default=[120, 160], nargs='*', type=int, help='learning rate schedule (when to drop lr by 10x)')# 动量更新值parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum of SGD solver')# 权重衰减速度parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float, metavar='W', help='weight decay (default: 1e-4)',dest='weight_decay')# 打印训练信息的频率parser.add_argument('-p', '--print-freq', default=10, type=int, metavar='N', help='print frequency (default: 10)')

◆ 代码说明：
 ① metavar 有助于 “生成帮助信息”，即方便我们对照着输入 “命令行参数”，我们可以不管。
 ② 默认情况下，解析器会将 “命令行参数” 当作简单字符串读入。 然而，“命令行字符串” 经常应当被解读为其他类型，例如 float 或 int。 add_argument() 的 type 关键字允许执行任何必要的类型检查和类型转换。如果 type 关键字使用了 default 关键字，则类型转换器仅会在默认值为字符串时被应用。
 ③ 有默认值，即包含有 default 的即是论文中找到的 “最佳参数值”。

# 模型地址(默认一开始没有), 如果你上次存了训练后的模型, 只需呀再利用这个 命令 写入 相应的地址, 即可继续加载该模型, 继续在其基础上训练parser.add_argument('--resume', default='', type=str, metavar='PATH', help='path to latest checkpoint (default: none)')# 还不太清楚...后期补充parser.add_argument('--world-size', default=-1, type=int, help='number of nodes for distributed training')# 还不太清楚...后期补充parser.add_argument('--rank', default=-1, type=int, help='node rank for distributed training')# 还不太清楚...后期补充parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str, help='url used to set up distributed training')# 分布式的后端parser.add_argument('--dist-backend', default='nccl', type=str, help='distributed backend')# 初始化训练的(随机)种子parser.add_argument('--seed', default=None, type=int, help='seed for initializing training. ')# 是否使用 GPU (默认有)parser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')# 使用多处理分布式训练启动, 每个节点有 N 个进程，即 N 个 gpu。无论是单节点还是多节点数据并行训练,这是使用 PyTorch 的最快方式。# action='store_true' 的意思是, 系统默认是 “启动的”parser.add_argument('--multiprocessing-distributed', action='store_true',help='Use multi-processing distributed training to launch ' 'N processes per node, which has N GPUs. This is the ' 'fastest way to use PyTorch for either single node or ' 'multi node data parallel training')

# MoCo v1 默认的配置：parser.add_argument('--moco-dim', default=128, type=int, help='feature dimension (default: 128)')parser.add_argument('--moco-k', default=65536, type=int, help='queue size; number of negative keys (default: 65536)')parser.add_argument('--moco-m', default=0.999, type=float, help='moco momentum of updating key encoder (default: 0.999)')parser.add_argument('--moco-t', default=0.07, type=float, help='softmax temperature (default: 0.07)')# MoCo v2 新增的配置：parser.add_argument('--mlp', action='store_true', help='use mlp head')parser.add_argument('--aug-plus', action='store_true', help='use moco v2 data augmentation')parser.add_argument('--cos', action='store_true', help='use cosine lr schedule')

2.3 关于“main()函数”的代码

def main():args = parser.parse_args()# 一般不会执行if args.seed is not None:# 如果要开启“系统随机”，这样会导致训练速度会大大下降random.seed(args.seed)torch.manual_seed(args.seed)cudnn.deterministic = Truewarnings.warn('You have chosen to seed training. ''This will turn on the CUDNN deterministic setting, ''which can slow down your training considerably! ''You may see unexpected behavior when restarting ''from checkpoints.')# 一般不会执行if args.gpu is not None:warnings.warn('You have chosen a specific GPU. This will completely ''disable data parallelism.')# 获取系统环境变量中 分布式训练 的值if args.dist_url == "env://" and args.world_size == -1:args.world_size = int(os.environ["WORLD_SIZE"])args.distributed = args.world_size > 1 or args.multiprocessing_distributedngpus_per_node = torch.cuda.device_count()# 获取 gpu 数量if args.multiprocessing_distributed:# Since we have ngpus_per_node processes per node, the total world_size needs to be adjusted accordingly# 由于每个节点都有 ngpus_per_node 个进程，因此需要相应地调整 world_size 的总数args.world_size = ngpus_per_node * args.world_size# Use torch.multiprocessing.spawn to launch distributed processes: the main_worker process function# 使用 torch.multiprocessing.spawn 启动分布式进程: main_worker 进程函数mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))else:# Simply call main_worker function# 如果不适用分布式训练, 那么就简单地调用 main_worker 进程函数main_worker(args.gpu, ngpus_per_node, args)

◆ 代码说明：
 ① 函数 parse_args(args=None, namespace=None) 的功能：将“参数字符串”(就是在运行该程序时，我们从键盘输入到这里面的)转换为 “对象” 并将其设为命名空间的属性。 返回带有成员的命名空间。
 ② mp.spawn() 第一个参数是一个函数，这个函数将执行训练的所有步骤。从这一步开始，Python 将建立多个进程，每个进程都会执行 main_worker() 函数。 第二个参数是开启的进程数目。第三个参数是传给 main_worker() 的函数实参。

2.4 关于“main_worker()函数”的代码

def main_worker(gpu, ngpus_per_node, args):args.gpu = gpu# suppress printing if not masterif args.multiprocessing_distributed and args.gpu != 0:def print_pass(*args):passbuiltins.print = print_passif args.gpu is not None:print("Use GPU: {} for training".format(args.gpu))if args.distributed:if args.dist_url == "env://" and args.rank == -1:args.rank = int(os.environ["RANK"])if args.multiprocessing_distributed:# For multiprocessing distributed training, rank needs to be the global rank among all the processes# 对于多处理分布式训练，RANK 必须是所有流程中的全局 RANKargs.rank = args.rank * ngpus_per_node + gpudist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,world_size=args.world_size, rank=args.rank)# create modelprint("=> creating model '{}'".format(args.arch))# 创建一个名为 “model” 的 MoCo 对象, 并传入一系列参数model = moco.builder.MoCo(models.__dict__[args.arch], args.moco_dim, args.moco_k, args.moco_m, args.moco_t, args.mlp)print(model) # 打印模型# 有 cuda 的话先查看是否打算用分布式计算(distributed computing)if args.distributed:# For multiprocessing distributed, DistributedDataParallel constructor# should always set the single device scope, otherwise,# DistributedDataParallel will use all available devices.# 对于分布式多处理，DistributedDataParallel 构造函数应始终设置单个设备范围, 否则 DistributedDataParallel 将使用所有可用设备。if args.gpu is not None:torch.cuda.set_device(args.gpu)model.cuda(args.gpu)# When using a single GPU per process and per# DistributedDataParallel, we need to divide the batch size# ourselves based on the total number of GPUs we have# 当每个进程和每个 DistributedDataParallel 使用单个 GPU 时, 我们需要根据拥有的 GPU 总数自行划分批大小(batch size)args.batch_size = int(args.batch_size / ngpus_per_node)args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])else: # 否则 DistributedDataParallel 将划分 batch_size 并将其分配给所有可用的 GPUmodel.cuda()# DistributedDataParallel will divide and allocate batch_size to all# available GPUs if device_ids are not setmodel = torch.nn.parallel.DistributedDataParallel(model)# 如果不打算采用分布式计算，且存在人为指定的gpuelif args.gpu is not None:torch.cuda.set_device(args.gpu)model = model.cuda(args.gpu)# comment out the following line for debuggingraise NotImplementedError("Only DistributedDataParallel is supported.")# raise: 用于抛出异常else:# AllGather implementation (batch shuffle, queue update, etc.) in this code only supports DistributedDataParallel.raise NotImplementedError("Only DistributedDataParallel is supported.")# raise: 用于抛出异常# define loss function (criterion) and optimizer# 定义损失函数(标准)和优化器criterion = nn.CrossEntropyLoss().cuda(args.gpu)optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=args.momentum, weight_decay=args.weight_decay)# optionally resume from a checkpoint# 如果可以从一个 “检查点” 恢复。“检查点”: 每一次迭代都会产生一个, 它即是训练模型的权重, 每隔一轮都会有所更新。if args.resume:if os.path.isfile(args.resume): # os.path.isfile(): 判断某一对象(需提供绝对路径)是否为文件print("=> loading checkpoint '{}'".format(args.resume))if args.gpu is None:# 如果没有指定 gpucheckpoint = torch.load(args.resume)else:# 如果指定了 gpu# Map model to be loaded to specified single gpu.# 加载模型到指定的单个 GPU 上loc = 'cuda:{}'.format(args.gpu)checkpoint = torch.load(args.resume, map_location=loc)# 从 checkpoint 里导出 epochargs.start_epoch = checkpoint['epoch']# 把 model 的 parameter, 即存有这些 parameter 的 'state_dict' 加载到当前 model 里来model.load_state_dict(checkpoint['state_dict']) # 把 optimizer 的信息, 即存有这些信息的 'optimizer' 加载到当前 optimizer 里来optimizer.load_state_dict(checkpoint['optimizer'])print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))else:print("=> no checkpoint found at '{}'".format(args.resume))# 该标志允许您启用内置的 cudnn 自动调谐器，以找到用于硬件的最佳算法。# 原理: 在 cudnn 中启用基准测试模式。 只要网络输入大小不变，基准测试模式就很好。 # 这样, cudnn 将为该特定配置寻找最佳算法集(这需要一些时间)。 这会使得运行时间加快。# 但是, 如果你的输入大小在每次迭代中都发生变化，则 cudnn 将在每次出现新大小时进行基准测试, 这可能会导致运行时性能变差。cudnn.benchmark = True# Data loading code# 加载数据traindir = os.path.join(args.data, 'train')# normalize 的初始化normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])# 如果使用 MoCo v2if args.aug_plus:# MoCo v2's aug: similar to SimCLR https://arxiv.org/abs/2002.05709augmentation = [transforms.RandomResizedCrop(224, scale=(0.2, 1.)),# 随机剪裁transforms.RandomApply([# 随机颜色抖动transforms.ColorJitter(0.4, 0.4, 0.4, 0.1)# not strengthened], p=0.8),transforms.RandomGrayscale(p=0.2),# 随机灰度transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=0.5),# 新加的数据增强手段:『高斯模糊』transforms.RandomHorizontalFlip(),# 随机翻转transforms.ToTensor(),normalize]else:# MoCo v1's aug: the same as InstDisc https://arxiv.org/abs/1805.01978augmentation = [transforms.RandomResizedCrop(224, scale=(0.2, 1.)),transforms.RandomGrayscale(p=0.2),transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),transforms.RandomHorizontalFlip(),transforms.ToTensor(),normalize]# moco.loader.TwoCropsTransform(): 将一个图像的两个随机剪裁作为“查询”和“键”。【该函数是个人写的, 在 loader.py 中】# transforms.Compose(): 将多个 transform 组合起来使用train_dataset = datasets.ImageFolder(traindir, moco.loader.TwoCropsTransform(transforms.Compose(augmentation)))# 有打算用分布式计算的话就定义 train_sampler 这个 object.# train_sampler 能在生成数据加载器(dataloader)时负责把 training data 划分到不同 gpu 上if args.distributed:train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)# 训练时采用的 “下采样” 手段else:train_sampler = None# 为了训练设置的 data_loader 对象(也可以说是一个变量)train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True)# 分布式训练开始for epoch in range(args.start_epoch, args.epochs):# 如果采用分布式训练，告诉 train_sampler 要训练多少个 epoch, 它好按周期来划分数据给到每次更新参数时用的 data_loaderif args.distributed:train_sampler.set_epoch(epoch) adjust_learning_rate(optimizer, epoch, args)# train for one epoch# 一次迭代训练train(train_loader, model, criterion, optimizer, epoch, args)# 调用“2.4的train()函数” if not args.multiprocessing_distributed or (args.multiprocessing_distributedand args.rank % ngpus_per_node == 0):save_checkpoint({# 每一次迭代都保存 “检查点”'epoch': epoch + 1,'arch': args.arch,'state_dict': model.state_dict(),'optimizer': optimizer.state_dict(),}, is_best=False, filename='checkpoint_{:04d}.pth.tar'.format(epoch))

◆ 补充代码说明：ToTensor()的功能：将矩阵『重塑(reshape)』为 (H, W, C) 的 “nump.ndarray 对象” 或将 img 『重塑(reshape)』转为 (C, H, W)的 “tensor 对象”，其将每一个数值归一化到 [0,1]，其归一化方法比较简单，直接除以 255 即可。

2.5 关于“train()函数”的代码

def train(train_loader, model, criterion, optimizer, epoch, args):batch_time = AverageMeter('Time', ':6.3f')# 创建一个名为 “Time” 的“AverageMeter对象”, 用于计算和存储平均值和当前值data_time = AverageMeter('Data', ':6.3f')# 6.3f 的意思是: 固定输出6个占位+显示到小数点后3位losses = AverageMeter('Loss', ':.4e')top1 = AverageMeter('Acc@1', ':6.2f')top5 = AverageMeter('Acc@5', ':6.2f')# 创建一个“ProgressMeter对象”, 用于实时输出训练的中间结果progress = ProgressMeter( len(train_loader), [batch_time, data_time, losses, top1, top5], prefix="Epoch: [{}]".format(epoch))# switch to train mode# 切换到模型的训练模式model.train()end = time.time()# 返回当前时间的时间戳(即从1970年1月1日至今经过的浮点秒数)# 从用来训练的dataloder里提取数据for i, (images, _) in enumerate(train_loader): # measure data loading time# 记录数据的加载时间data_time.update(time.time() - end) # 如果有指定 gpu，就把图片数据放入到指定gpu中，并采用非阻塞通信if args.gpu is not None:images[0] = images[0].cuda(args.gpu, non_blocking=True)images[1] = images[1].cuda(args.gpu, non_blocking=True)# compute output# 计算 q 与 k 的输出结果output, target = model(im_q=images[0], im_k=images[1])loss = criterion(output, target)# 计算对比损失# acc1/acc5 are (K+1)-way contrast classifier accuracy# measure accuracy and record loss# Top-1/Top-5是针对于 K 种分类的准确度, 并会记录“对比损失”acc1, acc5 = accuracy(output, target, topk=(1, 5))# 计算指定 k 值最高的准确度losses.update(loss.item(), images[0].size(0))# images.size(0)表示图片数量top1.update(acc1[0], images[0].size(0))top5.update(acc5[0], images[0].size(0))# compute gradient and do SGD step# 在 PyTorch 中, 我们需要在开始进行反向传播之前将梯度设置为零,# 因为 PyTorch 会在随后的向后传递中累积梯度。optimizer.zero_grad()# loss.backward(): 为每个具有 require_grad = True 的参数 x 计算 d(loss) / dx。 d(...)/dx是对“...”求导的意思# 这些对于每个参数 x 都累积到 x.grad 中。 伪代码: x.grad + = d(loss) / dxloss.backward()# optimizer.step(): 使用 x.grad 来更新 x 的值。 例如: SGD 优化器执行以下操作: x += -lr * x.grad optimizer.step()# measure elapsed time# 计算训练每一批次花费的时间batch_time.update(time.time() - end)end = time.time()# 重置 “批次开始训练的” 时间点# 按照某一设置好的 “打印频率” 来打印进度信息if i % args.print_freq == 0:progress.display(i)

◆ 补充代码说明：enumerate()的功能：对于一个可遍历的对象(如列表、字符串)，enumerate将其组成一个索引序列, 利用它可以同时获得索引和值。这个比用 range() 方便一点。

2.6 关于“AverageMeter类”的代码

class AverageMeter(object):"""Computes and stores the average and current value"""# 计算和存储平均值和当前值(关于准确率的), 这个是深度学习固定的模板def __init__(self, name, fmt=':f'):self.name = nameself.fmt = fmtself.reset()def reset(self):self.val = 0self.avg = 0self.sum = 0self.count = 0def update(self, val, n=1):self.val = valself.sum += val * nself.count += nself.avg = self.sum / self.countdef __str__(self):fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'return fmtstr.format(**self.__dict__)

◆ 补充代码说明：fmt 是 “format” 的缩写，用于格式化输出。

2.7 关于“ProgressMeter类”的代码

class ProgressMeter(object):# 这个也是深度学习固定的模板, 用于辅助“中间结果输出”def __init__(self, num_batches, meters, prefix=""):self.batch_fmtstr = self._get_batch_fmtstr(num_batches)self.meters = metersself.prefix = prefixdef display(self, batch):entries = [self.prefix + self.batch_fmtstr.format(batch)]entries += [str(meter) for meter in self.meters]print('\t'.join(entries))def _get_batch_fmtstr(self, num_batches):num_digits = len(str(num_batches // 1))fmt = '{:' + str(num_digits) + 'd}'return '[' + fmt + '/' + fmt.format(num_batches) + ']'

◆ 补充代码说明：a // b 的意思是对除以 b 的结果向负无穷方向取整后的数。例如：7 // 2 = 3。“prefix” 是用来输出时添加的前缀。

2.8 关于“adjust_learning_rate()函数”的代码

def adjust_learning_rate(optimizer, epoch, args):"""Decay the learning rate based on schedule"""# 根据训练进度递减学习速率# 这个也是深度学习固定的模板lr = args.lrif args.cos:# cosine lr schedulelr *= 0.5 * (1. + math.cos(math.pi * epoch / args.epochs))else:# stepwise lr schedulefor milestone in args.schedule:lr *= 0.1 if epoch >= milestone else 1.for param_group in optimizer.param_groups:param_group['lr'] = lr

2.9 关于“accuracy()函数”的代码

def accuracy(output, target, topk=(1,)):"""Computes the accuracy over the k top predictions for the specified values of k"""# 计算指定 k 值的最高准确度with torch.no_grad():# 不进行梯度下降maxk = max(topk)batch_size = target.size(0)_, pred = output.topk(maxk, 1, True, True)pred = pred.t()correct = pred.eq(target.view(1, -1).expand_as(pred))res = []for k in topk:correct_k = correct[:k].contiguous().view(-1).float().sum(0, keepdim=True)res.append(correct_k.mul_(100.0 / batch_size))return res

◆ 补充代码说明：原来版本的代码是没有 .contiguous() 的，加上才能跑。

2.10 最后一点

if __name__ == '__main__':main()

◆ 代码说明：这样写的目的是使得“文件只能作为脚本才能(直接)执行，而 import 到其他脚本中是不会被执行的。

三、builder.py —— MoCo模型建立的文件

3.1 关于“MoCo类”的代码

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reservedimport torchimport torch.nn as nnclass MoCo(nn.Module):"""Build a MoCo model with: a query encoder, a key encoder, and a queue使用查询编码器、键编码器和队列构建 MoCo 模型https://arxiv.org/abs/1911.05722"""def __init__(self, base_encoder, dim=128, K=65536, m=0.999, T=0.07, mlp=False):"""dim: feature dimension (default: 128)K: queue size; number of negative keys (default: 65536)m: moco momentum of updating key encoder (default: 0.999)T: softmax temperature (default: 0.07)dim: 特征维度(默认: 128)K: 查询字典的队列大小,也就是负样本的数量(默认: 65536)m: 键编码器更新的 Moco 动量值(默认: 0.999)T: softmax 的温度参数(默认: 0.07)"""# nn.Module 子类的函数必须在构建函数中执行父类的构造函数# 下式等价于 nn.Module.__init__(self)super(MoCo, self).__init__()self.K = Kself.m = mself.T = T# create the encoders# num_classes is the output fc dimension# Num_classes 是全连接层的输出维度self.encoder_q = base_encoder(num_classes=dim)self.encoder_k = base_encoder(num_classes=dim)if mlp:# hack: brute-force replacementdim_mlp = self.encoder_q.fc.weight.shape[1]# nn.Linear(): 从输入输出的张量的 shape 角度来理解, 相当于一个输入为[batch_size, in_features]的张量变换# 成了[batch_size, out_features]的输出张量。# nn.Sequential(): 可以将一系列的操作打包, 这些操作可以包括 Conv2d()、ReLU()、Maxpool2d()等, # 打包后方便调用吧，就相当于是一个黑箱，forward() 时调用这个黑箱就行了。self.encoder_q.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_q.fc)self.encoder_k.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), self.encoder_k.fc)# zip(): 该函数返回一个以元组为元素的列表，其中第 i 个元组包含每个参数序列的第 i 个元素。for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):param_k.data.copy_(param_q.data)# initialize. “键编码器”初始化时的参数直接拷贝“查询编码器”的param_k.requires_grad = False# not update by gradient. “键编码器”不进行反向传播更新权重# create the queue# self.register_buffer(‘my_buffer’, self.tensor)：my_buffer 是名字(str类型); self.tensor 是需要进行 register 登记的张量。# 这样我们就得到了一个新的张量, 这个张量会保存在 model.state_dict() 中，也就可以随着模型一起通过 .cuda() 复制到 gpu 上。self.register_buffer("queue", torch.randn(dim, K))# nn.functional.normalize(): 将某一个维度除以那个维度对应的范数self.queue = nn.functional.normalize(self.queue, dim=0)self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))# 使用装饰器 @torch.no_gard() 修饰的函数，在调用时不允许计算梯度@torch.no_grad()def _momentum_update_key_encoder(self):"""Momentum update of the key encoder动量键编码器的更新"""for param_q, param_k in zip(self.encoder_q.parameters(), self.encoder_k.parameters()):param_k.data = param_k.data * self.m + param_q.data * (1. - self.m)@torch.no_grad()def _dequeue_and_enqueue(self, keys):# 出队和入队操作# gather keys before updating queue# 在更新队列之前获取键keys = concat_all_gather(keys)# 个人写的函数batch_size = keys.shape[0]ptr = int(self.queue_ptr)# ptr 是指针(pointer)的缩写# assert: 检查条件，不符合就终止程序assert self.K % batch_size == 0# for simplicity# replace the keys at ptr (dequeue and enqueue)self.queue[:, ptr:ptr + batch_size] = keys.Tptr = (ptr + batch_size) % self.K# move pointerself.queue_ptr[0] = ptr@torch.no_grad()def _batch_shuffle_ddp(self, x):"""Batch shuffle, for making use of BatchNorm.*** Only support DistributedDataParallel (DDP) model. ***结合批量归一化进行“批量洗牌”, 仅支持分布式数据并行(DDP)计算模型。"""# gather from all gpus# 获取所有 gpu 的信息batch_size_this = x.shape[0]x_gather = concat_all_gather(x)batch_size_all = x_gather.shape[0]num_gpus = batch_size_all // batch_size_this# random shuffle index# 随机洗牌torch.randperm(): 将0~n-1(包括0和n-1)随机打乱后获得的数字序列，函数名是 random permutation 的缩写idx_shuffle = torch.randperm(batch_size_all).cuda() # broadcast to all gpus# 把洗好的序列重新传回所有 gputorch.distributed.broadcast(idx_shuffle, src=0)# index for restoring# torch.argsort(): 一般结合 torch.sort() 使用, 用于返回重排后的索引idx_unshuffle = torch.argsort(idx_shuffle)# shuffled index for this gpugpu_idx = torch.distributed.get_rank()# 返回当前进程组的排名idx_this = idx_shuffle.view(num_gpus, -1)[gpu_idx]return x_gather[idx_this], idx_unshuffle@torch.no_grad()def _batch_unshuffle_ddp(self, x, idx_unshuffle):"""Undo batch shuffle.*** Only support DistributedDataParallel (DDP) model. ***不做“批量洗牌”操作"""# gather from all gpusbatch_size_this = x.shape[0]x_gather = concat_all_gather(x)batch_size_all = x_gather.shape[0]num_gpus = batch_size_all // batch_size_this# restored index for this gpugpu_idx = torch.distributed.get_rank()idx_this = idx_unshuffle.view(num_gpus, -1)[gpu_idx]return x_gather[idx_this]# 模型的前向传播def forward(self, im_q, im_k):"""Input:im_q: a batch of query imagesim_k: a batch of key imagesOutput:logits, targets"""# compute query features# 计算“查询”的特征q = self.encoder_q(im_q)# queries: NxC, N 指有多少张图片, C 指有多少特征(维度)q = nn.functional.normalize(q, dim=1)# compute key features# 计算“键”的特征with torch.no_grad():# no gradient to keys. 将不用计算梯度的变量放在 with torch.no_grad() 里self._momentum_update_key_encoder()# update the key encoder, 更新“键”编码器# shuffle for making use of BNim_k, idx_unshuffle = self._batch_shuffle_ddp(im_k)k = self.encoder_k(im_k)# keys: NxCk = nn.functional.normalize(k, dim=1)# undo shufflek = self._batch_unshuffle_ddp(k, idx_unshuffle)# compute logits# Einstein sum is more intuitive# positive logits: Nx1l_pos = torch.einsum('nc,nc->n', [q, k]).unsqueeze(-1)# 计算“查询(query)”与“键,即正样本(key)”之间的相似度# negative logits: NxKl_neg = torch.einsum('nc,ck->nk', [q, self.queue.clone().detach()])#计算“查询(query)”与“负样本(来自队列Queue)”之间的相似度# logits: Nx(1+K)# torch.cat(): 是将两个张量(tensor)拼接在一起，cat 是 concatnate 的意思，即拼接，联系在一起。logits = torch.cat([l_pos, l_neg], dim=1)# apply temperature# 除以一个温度参数logits /= self.T# labels: positive key indicators# 建立一个 “标签矩阵”, 形状和 logits 一样, 矩阵中元素的数据类型为 torch.longlabels = torch.zeros(logits.shape[0], dtype=torch.long).cuda()# dequeue and enqueue# 出队和入队操作self._dequeue_and_enqueue(k)return logits, labels

3.2 关于“concat_all_gather()函数”的代码

# utils# 常用工具@torch.no_grad()def concat_all_gather(tensor):"""Performs all_gather operation on the provided tensors.*** Warning ***: torch.distributed.all_gather has no gradient."""# torch.ones_like(): 根据给定张量, 生成与其形状相同的全 1 张量(即矩阵).tensors_gather = [torch.ones_like(tensor)for _ in range(torch.distributed.get_world_size())]torch.distributed.all_gather(tensors_gather, tensor, async_op=False)output = torch.cat(tensors_gather, dim=0)return output

四、loader.py —— 辅助图片载入的文件

4.1 关于“TwoCropsTransform类”的代码

# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reservedfrom PIL import ImageFilterimport randomclass TwoCropsTransform:"""Take two random crops of one image as the query and key."""# 将一个图像的两个随机裁剪的图片作为“查询”和“键”def __init__(self, base_transform):self.base_transform = base_transformdef __call__(self, x):q = self.base_transform(x)k = self.base_transform(x)return [q, k]

4.2 关于“GaussianBlur类”的代码

class GaussianBlur(object):"""Gaussian blur augmentation in SimCLR https://arxiv.org/abs/2002.05709"""# 『高斯模糊增强(Gaussian blur augmentation)』的原理详见 SimCLRdef __init__(self, sigma=[.1, 2.]):self.sigma = sigmadef __call__(self, x):sigma = random.uniform(self.sigma[0], self.sigma[1])x = x.filter(ImageFilter.GaussianBlur(radius=sigma))return x

◆ 代码说明：MoCo v1 使用 RandomResizedCrop()，RandomGrayscale()，ColorJitter()，RandomHorizontalFlip() 这几种数据增强方式，到了 MoCo v2 又增加了：transforms.RandomApply([moco.loader.GaussianBlur([.1, 2.])], p=0.5) 这种方式。

五、main_lincls.py —— 评估模型建立的文件

● “lincls” 是 “Linear Classification” 的缩写。

● 整个评估模型的代码框架如下：

import ...parser.add_argument(...)def main():...def main_worker(gpu, ngpus_per_node, args):...def train(train_loader, model, criterion, optimizer, epoch, args):...def validate(val_loader, model, criterion, args):...def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):...def sanity_check(state_dict, pretrained_weights):"""Linear classifier should not change any weights other than the linear layer.This sanity check asserts nothing wrong happens (e.g., BN stats updated)."""...class AverageMeter(object):"""Computes and stores the average and current value"""def __init__(self, name, fmt=':f'):...def reset(self):...def update(self, val, n=1):...def __str__(self):...class ProgressMeter(object):def __init__(self, num_batches, meters, prefix=""):...def display(self, batch):...def _get_batch_fmtstr(self, num_batches):...def adjust_learning_rate(optimizer, epoch, args):"""Decay the learning rate based on schedule"""...def accuracy(output, target, topk=(1,)):"""Computes the accuracy over the k top predictions for the specified values of k"""...if __name__ == '__main__':main()

● Linear Evaluation 的代码整体结构和 Unsupervised Pre-training 是一致的，但是在做 Linear Evaluation 时要注意冻结 Encoder 的参数不更新，而只更新最后分类器的参数即可。代码表示如下：

def main_worker(gpu, ngpus_per_node, args):global best_acc1args.gpu = gpu...# create model...# freeze all layers but the last fc# 冻结 Encoder 的参数不更新，而只更新最后分类器的参数for name, param in model.named_parameters():if name not in ['fc.weight', 'fc.bias']:param.requires_grad = False# init the fc layer# 初始化分类器的参数model.fc.weight.data.normal_(mean=0.0, std=0.01)model.fc.bias.data.zero_()...# load from pre-trained, before DistributedDataParallel constructor# 下载预训练模型if args.pretrained:if os.path.isfile(args.pretrained):print("=> loading checkpoint '{}'".format(args.pretrained))checkpoint = torch.load(args.pretrained, map_location="cpu")# rename moco pre-trained keysstate_dict = checkpoint['state_dict']for k in list(state_dict.keys()):# retain only encoder_q up to before the embedding layer# 装入所有层之前保留“查询编码器”if k.startswith('module.encoder_q') and not k.startswith('module.encoder_q.fc'):# remove prefixstate_dict[k[len("module.encoder_q."):]] = state_dict[k]# delete renamed or unused kdel state_dict[k]args.start_epoch = 0msg = model.load_state_dict(state_dict, strict=False)assert set(msg.missing_keys) == {"fc.weight", "fc.bias"}print("=> loaded pre-trained model '{}'".format(args.pretrained))else:print("=> no checkpoint found at '{}'".format(args.pretrained))...# optimize only the linear classifier# 只优化线性分类器parameters = list(filter(lambda p: p.requires_grad, model.parameters()))assert len(parameters) == 2# fc.weight, fc.bias# 优化器 optimizer 作用的参数为 parameters, 它只包含分类器 fc 的 weight 和 bias 这两部分。optimizer = torch.optim.SGD(parameters, args.lr,momentum=args.momentum,weight_decay=args.weight_decay)...# Data loading code# 把训练集和测试集放进来traindir = os.path.join(args.data, 'train')valdir = os.path.join(args.data, 'val')...if args.evaluate:validate(val_loader, model, criterion, args)returnfor epoch in range(args.start_epoch, args.epochs):if args.distributed:train_sampler.set_epoch(epoch)adjust_learning_rate(optimizer, epoch, args)# train for one epoch# 训练一轮train(train_loader, model, criterion, optimizer, epoch, args)# evaluate on validation set# 在测试集上进行评估acc1 = validate(val_loader, model, criterion, args)# remember best acc@1 and save checkpoint# 保存 top-1 的准确度, 随后保存“寄放点”(代码在...中)is_best = acc1 > best_acc1best_acc1 = max(acc1, best_acc1)...

● 补充说明一：另外为了确保除了分类器 fc 的 weight 和 bias 的部分，其余所有参数不发生改变，作者使用了这个 sanity_check() 函数。这个函数有 2 个输入：state_dict( Linear Evaluation 进行了一个 Epoch 之后的模型 ) 和 pretrained_weights (预训练权重的存放文件夹)。先从 pretrained_weights 的位置导入权重，命名为 state_dict_pre。
 接下来通过 assert ((state_dict[k].cpu() == state_dict_pre[k_pre]).all()) 比较 state_dict 的权重和 pretrained_weights 是不是一致的，即 Encoder 的权重有没有发生变化。如果有变化就会打印 “k is changed in linear classifier training”。

● 补充说明二：“main_lincls.py” 比 “main_moco.py” 多了一个计算测试集的准确率的 validate() 函数，剩下的部分代码差不多一致的。

▶本文参考附录

MoCo v1原论文地址：https://arxiv.org/pdf/2003.04297.pdf.

MoCo v2原论文地址：https://arxiv.org/pdf/2003.04297.pdf.

[1] 《Self-Supervised Learning 超详细解读 (四)：MoCo系列解读 (1)》.

[2] 《ImageNet distributed training in PyTorch》.

[3] 《argparse — 命令行选项、参数和子命令解析器¶》.

⭐️ ⭐️