视频超分辨率#

在此示例中，我们使用 DALI 和 VideoReader 操作符为在 PyTorch 中实现的视频超分辨率网络训练提供数据。我们复制了 《End-to-End Learning of Video Super-Resolution with Motion Compensation》中描述的网络。

实现的网络 VSRNet 使用从视频中随机采样的奇数个连续帧作为输入。外部帧被扭曲以与中心帧对齐，使用的预训练光流网络为 FlowNetSD [2]。扭曲的帧然后被传递到卷积网络，该网络预测中心帧的单个更高分辨率版本。

我们对网络进行了一些小的修改，我们发现这些修改提高了收敛速度，例如批量归一化、Adam 优化器、循环学习率策略。

支持单节点单 GPU，支持 fp32 和 fp16 训练。

数据加载器#

提供了两个数据加载器选项以供比较

  • DALI 允许直接从 .mp4 文件随机访问帧序列，并使用 DALI 操作符进行增强

  • 标准 PyTorch 数据加载器从单个 .png 文件加载帧

这些数据加载器可以在 dataloading/dataloaders.py 中找到。

数据加载器性能#

我们展示了 DALI、NVVL 独立版和标准 .png 数据加载器在多种 VSRNet 训练场景中的性能特征。在所有情况下，我们都发现 DALI 降低了 CPU 负载、主机内存使用量、磁盘使用量和每次迭代数据加载时间。

数据加载器

输入分辨率

裁剪尺寸*

批量大小

数据集磁盘空间 (GB)

每次迭代数据时间 (毫秒)

DALI

720p

540p

3

0.961

0.28

NVVL

720p

540p

3

0.961

0.28

.png

720p

540p

3

38

2.66

  • 随机裁剪

要求#

用于运行此项目的系统必须包含两个或多个 Kepler、Pascal、Maxwell 或 Volta NVIDIA GPU。

软件要求

  • CUDA 9.0+

  • PyTorch 0.4+

  • FFmpeg=3.4.2

  • scikit-image

  • tensorflow

  • tensorboard

  • tensorboardX

FlowNet2-SD 实现和预训练模型#

我们使用了 此处提供的 FlowNet2-SD PyTorch 实现。它作为 git 子模块包含在此 repo 中。

为了使用预训练的 FlowNet2-SD 网络，请从该 repo 的根目录运行以下命令

git clone https://github.com/NVIDIA/flownet2-pytorch.git
pushd flownet2-pytorch
git checkout 6a0d9e70a5dcc37ef5577366a5163584fd7b4375
popd

训练此处实现的 VSRNet 需要使用来自 FlowNet2-SD 网络的预训练权重。我们提供了转换后的 Caffe 预训练模型，如下所示。如果您使用这些权重，请遵守许可协议

FlowNet2-SD[173MB]

训练代码将查找这些权重的默认位置是 flownet2-pytorch/networks/FlowNet2-SD_checkpoint.pth.tar。可以通过 main.py 的 --flownet_path 参数更改此位置。

数据#

按照 Makansi 等人的方法，我们使用 Myanmar 60p 视频作为原始数据源。

原始视频是 60 FPS、4K 分辨率的电影视频。为了准备用于训练的数据，您应该运行以下步骤

  1. 创建数据文件夹 <data_dir> 并下载 4K Myanmar 视频。

  2. 通过提供 4K Myanmar 视频文件作为参数来运行 prepare_data.sh。

这将在 <data_dir> 内部创建 5 个目录

  • orig

  • 540p

  • 720p

  • 1080p

  • 4k 这些目录中的每一个都包含场景（原始视频的短片段），分为 train 和 val 两个目录。

prepare_data.sh 中使用的脚本如下

tools/split_scenes.py：将视频分割成场景并删除音轨

python ./tools/split_scenes.py --raw_data <path_to_mp4_file> --out_data <data_dir>

场景将被写入 <data_dir>/orig/scenes。场景将被分为训练和验证文件夹。

tools/transcode_scenes.py：转码场景以获得更小的关键帧间隔，并可能降低分辨率

python ./tools/transcode_scenes.py --main_data <data_dir> --resolution <resolution>

其中 <resolution> 可以是以下之一：“4K”、1080p、720p 或 540p。转码后的场景将被写入 <data_dir>/<resolution>/scenes 并分为训练和验证文件夹。运行带有 –help 的脚本以查看更多选项。请注意，虽然您可以一步分割和转码原始视频，但我们发现先分割再转码要快得多。

训练#

可以通过运行 docs/examples/use_cases/video_superres/run.sh 中提供的以下命令来运行训练

./run.sh

此文件允许配置各种训练选项 - 预计您将根据您的系统适当修改数据路径。

训练数据的可视化，例如损失曲线和时间，以及示例图像，通过 Tensorboard 以及 tensorboardX 库提供。在训练运行时，您可以在 <host_ip>:6006 访问 Tensorboard。

Myanmar 验证集上的结果#

此项目的所有测试均在 NVIDIA DGX-1 上使用所有 8 个 V100 GPU 进行，运行 CUDA 9.1、PyTorch 0.4.0a0+02b758f、cuDNN v7.0.5，在 Ubuntu 16.04 Docker 容器中。

输入图像 (128x240 - 点击查看实际尺寸)

VSRNet 预测 (512x960 - 点击查看实际尺寸)

示例训练损失 (fp16, 批量大小 7, min_lr=max_lr=0.001)

示例验证 PSNR (fp16, 批量大小 7, min_lr=max_lr=0.001)

参考#

如果您发现此实现对您的工作有用，请适当致谢并引用以下论文

@InProceedings{IB17,
  author       = "O. Makansi and E. Ilg and and Thomas Brox",
  title        = "End-to-End Learning of Video Super-Resolution with Motion
                  Compensation",
  booktitle    = "German Conference on Pattern Recognition (GCPR) 2017",
  month        = " ",
  year         = "2017",
  url          = "http://lmb.informatik.uni-freiburg.de/Publications/2017/IB17"
}

@InProceedings{IMKDB17,
  author       = "E. Ilg and N. Mayer and T. Saikia and M. Keuper and A.
                  Dosovitskiy and T. Brox",
  title        = "FlowNet 2.0: Evolution of Optical Flow Estimation with
                  Deep Networks",
  booktitle    = "IEEE Conference on Computer Vision and Pattern Recognition
                  (CVPR)",
  month        = "Jul",
  year         = "2017",
  url          = "http://lmb.informatik.uni-freiburg.de//Publications/2017/IMKDB17"
}

此示例和此 README 基于 NVVL 的 PyTorch 示例。