用于填充和图像修复的局部卷积层

填充论文 | 修复论文 | 修复 YouTube 视频 | 在线修复演示

这是局部卷积层的 PyTorch 实现。它可以作为一种新的填充方案，也可以用于图像修复。

基于局部卷积的填充
Guilin Liu、Kevin J. Shih、Ting-Chun Wang、Fitsum A. Reda、Karan Sapra、Zhiding Yu、Andrew Tao、Bryan Catanzaro
NVIDIA 公司
技术报告 (Technical Report) 2018

使用局部卷积修复不规则孔洞的图像
Guilin Liu、Fitsum A. Reda、Kevin J. Shih、Ting-Chun Wang、Andrew Tao、Bryan Catanzaro
NVIDIA 公司
2018 年欧洲计算机视觉大会 (ECCV)

与零填充的对比

安装

安装说明请参见：https://github.com/pytorch/examples/tree/master/imagenet

使用方法：

• 使用局部卷积进行填充

# 带零填充的标准卷积层
nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)

# 基于局部卷积的填充
PartialConv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)

• 使用局部卷积进行图像修复时，将 multi_channel 和 return_mask 都设置为 True

# 用于修复的局部卷积（使用多通道并更新掩码）
PartialConv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False, multi_channel=True, return_mask=True)

使用 AMP 进行图像修复的混合精度训练

• 安装：要启用混合精度支持，请先从以下地址安装 apex：https://github.com/NVIDIA/apex
• 必需更改 #1（常规更改）：AMP 所需的常规更改

from apex import amp

# 初始化模型和优化器
self.model, self.optimizer = amp.initialize(self.model, self.optimizer, opt_level=args.amp_opt_level)

# 初始化 VGG 损失函数/提取器
self.vgg_feat_loss = amp.initialize(self.vgg_feat_loss, opt_level=args.amp_opt_level)

# 缩放损失
with amp.scale_loss(total_loss, self.g_optimizer) as scaled_loss:
    scaled_loss.backward()

• 必需更改 #2（Gram 矩阵损失）：在计算 Gram 矩阵损失时，将一步除法改为两步小除法
  • 从 一步除法：gram = features.bmm(features_t) / (ch * h * w)
  • 改为 两步小除法：

input = torch.zeros(b, ch, ch).type(features.type())
gram = torch.baddbmm(input, features, features_t, beta=0, alpha=1./(ch * h * w), out=None)

• 必需更改 #3（小常数）：将小常数稍微增大（例如从 1e-8 增至 1e-6）
  • 从 1e-8：self.mask_ratio = self.slide_winsize/(self.update_mask + 1e-8)
  • 至 1e-6：self.mask_ratio = self.slide_winsize/(self.update_mask + 1e-6)

使用局部卷积填充训练 ImageNet

• 使用零填充的 ResNet50（默认填充）

python main.py -a resnet50 --data_train /path/ILSVRC/Data/CLS-LOC/train --data_val /path/ILSVRC/Data/CLS-LOC/perfolder_val --batch-size 192 --workers 32 --prefix multigpu_b192 --ckptdirprefix experiment_1/

• 使用局部卷积填充的 ResNet50

python main.py -a pdresnet50 --data_train /path/ILSVRC/Data/CLS-LOC/train --data_val /path/ILSVRC/Data/CLS-LOC/perfolder_val --batch-size 192 --workers 32 --prefix multigpu_b192 --ckptdirprefix experiment_1/

• 使用零填充的 vgg16_bn（默认填充）

python main.py -a vgg16_bn --data_train /path/ILSVRC/Data/CLS-LOC/train --data_val /path/ILSVRC/Data/CLS-LOC/perfolder_val --batch-size 192 --workers 32 --prefix multigpu_b192 --ckptdirprefix experiment_1/

• 使用局部卷积填充的 vgg16_bn

python main.py -a pdvgg16_bn --data_train /path/ILSVRC/Data/CLS-LOC/train --data_val /path/ILSVRC/Data/CLS-LOC/perfolder_val --batch-size 192 --workers 32 --prefix multigpu_b192 --ckptdirprefix experiment_1/

带有局部卷积填充的 VGG 和 ResNet 网络的预训练检查点（权重）：

https://www.dropbox.com/sh/t6flbuoipyzqid8/AACJ8rtrF6V5b9348aG5PIhia?dl=0

五次运行中与零填充、反射填充和复制填充的对比

每次运行采用单裁剪测试时的最佳 top-1 准确率。*_zero、*_pd、*_ref 和 *_rep 分别表示对应模型分别采用零填充、局部卷积填充、反射填充和复制填充。*_best 表示每次训练运行中的最佳验证分数。平均值代表 5 次运行的平均准确率。差异列表示与采用零填充的相应网络相比的差异。标准差列表示 5 次运行准确率的标准差。PT_official 表示 PyTorch 官网发布的相应官方准确率：https://pytorch.org/docs/stable/torchvision/models.html

引用

@inproceedings{liu2018partialpadding,
   author    = {刘桂林、凯文·J·施、王廷春、菲茨姆·A·雷达、卡兰·萨普拉、于志东、安德鲁·陶、布莱恩·卡坦扎罗},
   title     = {基于部分卷积的填充},
   booktitle = {arXiv 预印本 arXiv:1811.11718},   
   year      = {2018},
}
@inproceedings{liu2018partialinpainting,
   author    = {刘桂林、菲茨姆·A·雷达、凯文·J·施、王廷春、安德鲁·陶、布莱恩·卡坦扎罗},
   title     = {使用部分卷积对不规则孔洞进行图像修复},
   booktitle = {欧洲计算机视觉大会 (ECCV)},   
   year      = {2018},
}

联系人：刘桂林（guilinl@nvidia.com）

致谢

我们感谢顾金伟、马蒂厄·勒、安杰伊·苏莱基、马雷克·科沃杰和刘洪福的有益讨论。

PartialConv 快速上手指南

PartialConv 是 NVIDIA 开源的 PyTorch 实现，主要包含偏卷积层（Partial Convolution Layer）。它既可作为优于传统零填充（Zero Padding）的新型填充方案用于图像分类任务，也可专门用于处理不规则缺失区域的**图像修复（Image Inpainting）**任务。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python: 3.6+
• 深度学习框架: PyTorch (建议最新稳定版)
• GPU: NVIDIA GPU (可选，但强烈推荐用于训练和推理)
• 前置依赖:
  • 基础 PyTorch 环境安装参考：PyTorch Examples - ImageNet
  • 混合精度训练支持 (可选): 若需进行图像修复的混合精度训练，需预先安装 NVIDIA Apex：

    git clone https://github.com/NVIDIA/apex
    cd apex
    pip install -v --no-cache-dir ./
    

    国内加速提示: 若下载缓慢，可使用清华源或阿里源安装依赖： pip install -v --no-cache-dir ./ -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

本项目主要为 PyTorch 模块实现，通常无需复杂的编译安装过程。您可以直接克隆仓库并将 models 目录下的代码集成到您的项目中，或者将其作为子模块引用。

git clone https://github.com/NVIDIA/partialconv.git
cd partialconv

注：核心逻辑位于 models/partialconv2d.py，使用时只需将该文件或其类定义引入您的项目即可。

基本使用

1. 作为新型填充方案 (Padding)

在传统卷积中替换 nn.Conv2d 即可实现基于偏卷积的填充，通常能获得比零填充更好的特征提取效果。

from models.partialconv2d import PartialConv2d
import torch.nn as nn

# 传统卷积层 (零填充)
# nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)

# 替换为偏卷积填充层
model_layer = PartialConv2d(3, 16, kernel_size=3, stride=1, padding=1, bias=False)

2. 用于图像修复 (Image Inpainting)

若用于图像修复任务，必须将 multi_channel 和 return_mask 参数设置为 True，以便网络能够处理多通道输入并动态更新掩码（Mask）。

from models.partialconv2d import PartialConv2d

# 配置用于图像修复的偏卷积层
inpainting_layer = PartialConv2d(
    3, 16, 
    kernel_size=3, 
    stride=1, 
    padding=1, 
    bias=False, 
    multi_channel=True, 
    return_mask=True
)

3. 训练示例 (ImageNet 分类)

如果您希望复现论文中在 ImageNet 上的分类结果，可以使用提供的 main.py 脚本。

使用 ResNet50 + 偏卷积填充：

python main.py -a pdresnet50 --data_train /path/ILSVRC/Data/CLS-LOC/train --data_val /path/ILSVRC/Data/CLS-LOC/perfolder_val --batch-size 192 --workers 32 --prefix multigpu_b192 --ckptdirprefix experiment_1/

使用 VGG16_bn + 偏卷积填充：

python main.py -a pdvgg16_bn --data_train /path/ILSVRC/Data/CLS-LOC/train --data_val /path/ILSVRC/Data/CLS-LOC/perfolder_val --batch-size 192 --workers 32 --prefix multigpu_b192 --ckptdirprefix experiment_1/

预训练模型: 作者提供了基于偏卷积填充的 VGG 和 ResNet 预训练权重，可从 Dropbox 链接 下载。