DCN V1代码阅读笔记

最新推荐文章于 2024-07-06 03:29:01 发布

原创

最新推荐文章于 2024-07-06 03:29:01 发布 · 1.4k 阅读

15 ·

CC 4.0 BY-SA版权

1. 前言

笔者前几天阅读了MASA 的可变形卷积算法，并写了一篇算法笔记：MASA DCN(可变形卷积) 算法笔记，然后里面也说了会放出一个代码解读出来，所以今天的推文就是来干这件事的，希望看完这篇文章你可对DCN的细节可以更加了解。本次解析的代码来自：https://github.com/ChunhuanLin/deform_conv_pytorch 。

2. 代码整体介绍

打开这个工程，我们发现只有 $3$ 个文件，整个结构非常简单：

代码结构
实际上我们关注前面 $2$ 个文件deform_conv.py和demo.py就可以了，至于test_against_mxnet.ipynb只是测试了以下mxnet自带的可变形卷积op和这个库里面实现的可变形卷积op对同一个输入进行处理，输出结果是否一致，从jupyter的结果来看是完全一致的，也证明作者这个DeformConv2D是正确实现了，接下来我们就从源码角度来看一下可变形卷积究竟是如何实现的。

另外需要注意的是这个Pytorch代码版本是0.4.1，如果需要使用1.0以上的可能需要稍微改改某些API，但改动应该很少的，毕竟这个代码加起来也就2,300行。。

3. 可变形卷积示意图

为了更好的理解代码，我们先形象的回忆一下可变形卷积。

下面的Figure2展示了可变形卷积的示意图：

Deformable Convolution的示意图

可以看到可变形卷积的结构可以分为上下两个部分，上面那部分是基于输入的特征图生成offset，而下面那部分是基于特征图和offset通过可变形卷积获得输出特征图。

假设输入的特征图宽高分别为 $w$ ， $h$ ，下面那部分的卷积核尺寸是 $k_h$ 和 $k_w$ ，那么上面那部分卷积层的卷积核数量应该是 $2\times k_h\times k_w$ ，其中 $2$ 代表 $x$ ， $y$ 两个方向的offset。

并且，这里输出特征图的维度和输入特征图的维度一样，那么offset的维度就是 $2\times k_h\times k_w, h, w]$ ，假设下面那部分设置了group参数（代码实现中默认为 $4$ ），那么第一部分的卷积核数量就是 $2\times k_h \times k_w \times group$ ，即每一个group共用一套offset。下面的可变形卷积可以看作先基于上面那部分生成的offset做了一个插值操作，然后再执行普通的卷积。

4. 解析`demo.py`

demo.py是训练和参数设置的入口，我们进来看看。

首先是设置一些训练参数，并加载MNIST数据集到train_loader和test_loader里面：

# Training settings
parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
# 设置训练batch_size，默认为32
parser.add_argument('--batch-size', type=int, default=32, metavar='N',
                    help='input batch size for training (default: 32)')
# 设置测试的batch_size，默认为32
parser.add_argument('--test-batch-size', type=int, default=32, metavar='N',
                    help='input batch size for testing (default: 32)')
# 设置训练epochs数，默认为10
parser.add_argument('--epochs', type=int, default=10, metavar='N',
                    help='number of epochs to train (default: 10)')
# 设置学习率，默认为0.01
parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                    help='learning rate (default: 0.01)')
# 设置SGD的动量参数，默认为0.5
parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                    help='SGD momentum (default: 0.5)')
# 设置是否使用GPU训练
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='disables CUDA training')
# 设置随机数种子
parser.add_argument('--seed', type=int, default=1, metavar='S',
                    help='random seed (default: 1)')
# 没用到这个参数，不用管
parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                    help='how many batches to wait before logging training status')

# 解析参数
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()

torch.manual_seed(args.seed)
if args.cuda:
    # 为当前GPU设置随机种子
    # 如果使用多个GPU，应该使用torch.cuda.manual_seed_all()为所有的GPU设置种子
    torch.cuda.manual_seed(args.seed)

kwargs = {
   
   'num_workers': 1, 'pin_memory': True} if args.cuda else {
   
   }

# 加载数据。组合数据集和采样器，提供数据上的单或多进程迭代器
# 参数：
# dataset：Dataset类型，从其中加载数据
# batch_size：int，可选。每个batch加载多少样本
# shuffle：bool，可选。为True时表示每个epoch都对数据进行洗牌
# sampler：Sampler，可选。从数据集中采样样本的方法。
# num_workers：int，可选。加载数据时使用多少子进程。默认值为0，表示在主进程中加载数据。
# collate_fn：callable，可选。
# pin_memory：bool，可选
# drop_last：bool，可选。True表示如果最后剩下不完全的batch,丢弃。False表示不丢弃。

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./MNIST', train=True, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=args.batch_size, shuffle=True, **kwargs)
test_loader = torch.utils.data.DataLoader(
    datasets.MNIST('./MNIST', train=False, download=True,
                   transform=transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])),
    batch_size=args.test_batch_size, shuffle=True, **