TensorFlow 公式の CycleGAN チュートリアルをオリジナルデータで動かす

公式の CycleGAN チュートリアルをオリジナルデータで動かした時のメモです。Google Colab 上で動かします。

1. データの準備

リンクがある cyclegan.ipynb はデフォルトで tensorflow_datasets の cycle_gan/horse2zebra を利用するようになっています。順当にやるなら tfds new などでカスタムデータセットを作成する形なります。
ただ若干書くことがありそうだったので、以下に書いたように horse2zebra.zip の DL を置換する方法で動かしました。

1.1. horse2zebra.zip の DL を置換

cycle_gan/horse2zebra の実装は大学のサイトから horse2zebra.zip を落としてくる実装になっています。
ただ試すだけであれば、horse2zebra.zip を落としてくる部分を置換すれば良いので、まずデータを以下のように作成します。

• horse2zebra
  • testA
    • testA_0001.jpg
    • testA_0002.jpg
    • ...
  • testB
    • testB_0001.jpg
    • testB_0002.jpg
    • ...
  • trainA
    • trainA_0001.jpg
    • trainA_0002.jpg
    • ...
  • trainB
    • trainB_0001.jpg
    • trainB_0002.jpg
    • ...

上記の構成 (horse2zebra 含め) を zip を固めて Google Drive に上げておき(例: my_horse2zebra.zip)、以下のコードを cyclegan.ipynb の冒頭にいれると、ダウンロードの代わりに zip を利用できます。

## Google Drive をマウント
from google.colab import drive
drive.mount('/content/drive')

## Google Drive から zip をコピー
!cp /content/drive/My\ Drive/my_horse2zebra.zip ./

## URL をダウンロードするところを extract に変更
## https://www.tensorflow.org/datasets/api_docs/python/tfds/download/DownloadManager#extract
!git clone --depth 1 https://github.com/tensorflow/datasets
!perl -i.org -pe 's{data_dirs = .*}{data_dirs = dl_manager.extract("/content/my_horse2zebra.zip")};' datasets/tensorflow_datasets/image_classification/cycle_gan.py
!diff datasets/tensorflow_datasets/image_classification/cycle_gan.py{.org,}
!cd datasets/; pip install -e .

## import しておく
import sys
sys.path.append("/content/datasets/")
import tensorflow_datasets as tfds
print(tfds.__version__)

2. 学習

「Checkpoints」では checkpoint_path に前回のチェックポイントがあれば自動的に読み込んでくれるので (デフォルトだと ./checkpoints/train = /content/checkpoints/train)、Google Drive 上のディレクトリのシムリンクを作っておきます。

!mkdir -p /content/drive/My\ Drive/cyclegan_checkpoints/
!ln -s /content/drive/My\ Drive/cyclegan_checkpoints/ /content/checkpoints

後は cyclegan.ipynb をただ上から実行すれば学習・出力されます。