SIGGRAPH Asia2021論文のTime Travel Rephotographyを試してみた


こちらの論文のソースコードが公開されているので動かしてみます。
https://github.com/Time-Travel-Rephotography/Time-Travel-Rephotography.github.io

元のソースコードではtorch==1.4.0+cu100となっているのですが、手元のRTX3080Tiではcuda11以降しか使えないのでPyTorchのバージョンを上げています。RTX2000系までの環境なら元のソースコードのrequirements.txtで動くと思います。


##動作環境(ホストマシン)
- Ubuntu20.04
- RTX3080Ti

ホスト側のnvidia-smiの結果

```
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 510.47.03    Driver Version: 510.47.03    CUDA Version: 11.6     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  Off  | 00000000:01:00.0 Off |                  N/A |
|  0%   42C    P8    27W / 400W |     68MiB / 12288MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

```

##動作環境(docker)
- docker image: nvidia/cuda:11.3.0-devel-ubuntu20.04
- cuda 11.3 (RTX3000系なのでcuda11以降が必要)
- python3.8.8
- 環境変数: MKL_THREADING_LAYER=GNU


##環境変数
環境変数のMKL_THREADING_LAYERとMKL_SERVICE_FORCE_INTELまわりのエラーが出たので設定しました。MKL_SERVICE_FORCE_INTEL=1で回避できるという記事もあったのですがダメでした。MKL_THREADING_LAYER=GNUで回避できたようです。ここは詳しく調べてないので適当です。
https://github.com/pytorch/pytorch/issues/37377

##PyTorch
PyTorchはStableではなくNightlyのほうを入れています。Stableでは最新のcudaに対応していないためです。

```
RUN pip install --pre torch torchvision torchaudio -f https://download.pytorch.org/whl/nightly/cu113/torch_nightly.html
```

##それ以外のモジュール(requirements.txt)

```
numpy
Pillow
ninja
tqdm
opencv-python
scikit-image
tensorboard
tensorflow
dlib
scipy
matplotlib
pprintpp
gdown
```


チェックポイントをダウンロードします。公式のドキュメントはファイル名が間違っています。

```
./scripts/download_checkpoints.sh
```

以上の設定で動きました。結果はresults以下に出力されます。

```
./scripts/run.sh b "dataset/Abraham Lincoln_01.png" 0.75 
```

実行結果

```
root@host_time_travel:/data_root/Time-Travel-Rephotography.github.io# ./scripts/run.sh b "dataset/Abraham Lincoln_01.png" 0.75
++ spectral_sensitivity=b
++ path='dataset/Abraham Lincoln_01.png'
++ blur_radius=0.75
+++ dirname 'dataset/Abraham Lincoln_01.png'
++ list=dataset
+++ basename dataset
++ list=dataset
++ '[' b == b ']'
++ FLAGS=(--spectral_sensitivity b --encoder_ckpt checkpoint/encoder/checkpoint_b.pt)
++ name='dataset/Abraham Lincoln_01'
++ name='Abraham Lincoln_01'
++ echo 'Abraham Lincoln_01'
Abraham Lincoln_01
++ python projector.py 'dataset/Abraham Lincoln_01.png' --gaussian 0.75 --log_dir log/ --results_dir results/ --spectral_sensitivity b --encoder_ckpt checkpoint/encoder/checkpoint_b.pt
------------ Parameters -------------
camera_lr: 0.01
ckpt: checkpoint/stylegan2-ffhq-config-f.pt
coarse_min: 32
color_transfer: 10000000000.0
contextual: 0.1
cx_layers: ['relu3_4', 'relu2_2', 'relu1_2']
e4e_ckpt: checkpoint/e4e_ffhq_encode.pt
e4e_size: 256
encoder_ckpt: checkpoint/encoder/checkpoint_b.pt
encoder_size: 256
eye: 0.1
gaussian: 0.75
generator_size: 1024
init_latent: None
input: dataset/Abraham Lincoln_01.png
log_dir: log/
log_freq: 10
log_visual_freq: 50
lr: 0.1
mix_layer_range: [10, 18]
noise_ramp: 0.75
noise_regularize: 50000.0
noise_strength: 0.0
rand_seed: None
recon_size: 256
results_dir: results/
spectral_sensitivity: b
vgg: 1
vggface: 0.3
wplus_step: [250, 750]
-------------------------------------
b-G0.75-init(10,18)-s256-vgg1-vggface0.3-eye0.1-color1.0e+10-cx0.1(relu3_4,relu2_2,relu1_2)-NR5.0e+04-lr0.1_0.01-c32-wp(250,750)
Loading e4e over the pSp framework from checkpoint: checkpoint/e4e_ffhq_encode.pt
python test.py /data_root/Time-Travel-Rephotography.github.io/results/Abraham Lincoln_01/input_sibling /data_root/Time-Travel-Rephotography.github.io/results/Abraham Lincoln_01/skin_mask/face_parsing
Downloading: "https://download.pytorch.org/models/resnet18-5c106cde.pth" to /root/.cache/torch/hub/checkpoints/resnet18-5c106cde.pth
100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 44.7M/44.7M [00:00<00:00, 112MB/s]
[ 0  1  2  3  4  7  8 10 12 13 14 16 17]
[ 0  1  2  3  4  5  7  8 10 12 13 14 16 17]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:00<00:00, 357.42it/s]
/root/.pyenv/versions/anaconda3-2021.05/lib/python3.8/site-packages/torch/functional.py:568: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2292.)
  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]
Downloading: "https://download.pytorch.org/models/vgg16-397923af.pth" to /root/.cache/torch/hub/checkpoints/vgg16-397923af.pth
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 528M/528M [00:05<00:00, 108MB/s]
Downloading: "https://download.pytorch.org/models/vgg19-dcbb9e9d.pth" to /root/.cache/torch/hub/checkpoints/vgg19-dcbb9e9d.pth
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 548M/548M [00:04<00:00, 116MB/s]
Optimizing 32x32
  0%|                                                                                                                                   | 0/250 [00:00<?, ?it/s]/data_root/Time-Travel-Rephotography.github.io/optim/radam.py:60: UserWarning: This overload of addcmul_ is deprecated:
        addcmul_(Number value, Tensor tensor1, Tensor tensor2)
Consider using one of the following signatures instead:
        addcmul_(Tensor tensor1, Tensor tensor2, *, Number value) (Triggered internally at  ../torch/csrc/utils/python_arg_parser.cpp:1055.)
  exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
face:  3.603e+00; eye:  4.430e+00; contextual:  7.698e+00; color_transfer:  0.000e+00; noise:  1.715e-01:   0%|                         | 0/250 [00:00<?, ?it/s]/root/.pyenv/versions/anaconda3-2021.05/lib/python3.8/site-packages/torchvision/utils.py:63: UserWarning: The parameter 'range' is deprecated since 0.12 and will be removed in 0.14. Please use 'value_range' instead.
  warnings.warn(
face:  2.316e+00; eye:  3.881e+00; contextual:  8.728e+00; color_transfer:  3.332e-13; noise:  7.538e-04: 100%|███████████████| 250/250 [01:01<00:00,  4.07it/s]
Optimizing 64x64
face:  1.667e+00; eye:  3.558e+00; contextual:  7.689e+00; color_transfer:  4.897e-13; noise:  5.377e-05: 100%|███████████████| 750/750 [03:04<00:00,  4.07it/s]

real    5m10.798s
user    5m20.279s
sys     0m16.268s
```


![Abraham Lincoln_01-b-G0.75-init(10,18)-s256-vgg1-vggface0.3-eye0.1-color1.0e+10-cx0.1(relu3_4,relu2_2,relu1_2)-NR5.0e+04-lr0.1_0.01-c32-wp(250,750) (1).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/935724/6860755f-c41c-5fe6-571c-e8a61b06f466.png)