Pytorch+Tensorflowのちゃんぽんコードのすゝめ（tfdsでpytorchをブーストさせる話）

この記事は、
『Pytorchユーザーが実験速度で悩んでるときに、Tensorflowを部分的に使ってみるの割とアリでは説』
をちょっと検証してみたという内容になります

# とどのつまり(編集 2020/07/20)
ちゃんぽんの方が純粋pytorchより速かったりもしますが、いやでもやっぱ普通はフレームワーク統一した方がいいですたぶん。(可読性、コードの再利用性的に)

# 背景(ちょっとポエムを…本題は『やったこと』からです)

## Pytorch v.s. Tensorflow

先日こんなツイートを見かけました(引用失礼します:bow:)

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">TensorFlowが書きにくいというのはどういうことだろう．インターフェイスとして書きにくいというのであれば，それはないと思う，というかどのツールも正直大して変わらない．エラーの出処が深すぎて（あとデコレータでスタック改造していたりするので）デバッグできないというのは分かる．</p>&mdash; Yusuke Oda (@odashi_t) <a href="https://twitter.com/odashi_t/status/1202588933952983041?ref_src=twsrc%5Etfw">December 5, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">TF(v1)だと一度走れば完璧に動くことが期待できるのでデプロイに強いのと，いくらでもグラフの最適化の余地があるあたりが長所で，一方でこういうのは研究向けには割とどうでも良いので，実際にTFが製品・PyTorchが研究方面でよく使われているのを見ると棲み分けは出来ている気がしますね．</p>&mdash; Yusuke Oda (@odashi_t) <a href="https://twitter.com/odashi_t/status/1202593781385547781?ref_src=twsrc%5Etfw">December 5, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

<blockquote class="twitter-tweet"><p lang="ja" dir="ltr">ただ、企業での利用において、PyTorchでは困難なほどのグラフの最適化が要求される割合が、今後（金銭的コスト当たりの計算力や、計算力あたりの精度の向上なども考慮したうえで）どの程度変化するかで、ニーズが大きく変わる可能性もあると思っています。</p>&mdash; じんべえざめ (@jinbeizame007) <a href="https://twitter.com/jinbeizame007/status/1202598637164888064?ref_src=twsrc%5Etfw">December 5, 2019</a></blockquote> <script async src="https://platform.twitter.com/widgets.js" charset="utf-8"></script>

個人的にまとめると↓こんな感じだと思ってます

- Tensorflowは速度・デプロイ重視でPytorchはデバッグしやすさ重視??（辿ればおそらくDefine-and|by-Runの話）
- ただ重たい実験はPytorchだと渋いかも…
- そんなことより、国産派はPytorch一択。(?)（ChainerがPytorchに合流することになったそうですし…）
- そんなそんなことより、**深層生成モデルを[Pixyz](https://github.com/masa-su/pixyz)で、深層強化学習を[Machina](https://github.com/DeepX-inc/machina)で書きたい**ので、問答無用でPytorch

まぁぶっちゃけ一番最後ですね！！！

## Pytorch高速化の日々
- 完全に余談
  - 著者はあるメタ学習手法の再現をPytorchで書いたもののクソ遅かったので、怒涛の並列化で10倍速くらいにした経験があります。これはたぶんうまくいきました（元が遅かった:disappointed_relieved:）。
  - 著者は別のメタ学習手法の再現をPytorchで書いたもののクソ遅くて実装もバグっていたときに、高速化ばっかして速くはなったものの結局バグったまま精度が出ず締め切りに間に合わなかった経験があります。（順番を考えましょう…:disappointed_relieved:）
- はいつらい。はい苦い。

## Pytorchがなぜ(比較的?)遅いのか
- Define-and-Run思想の設計ゆえだと思うんですが、Tensorflowはepoch回す前の最適化が強そう
- というかPytorchは特にデータの読み込みとか前処理が遅そう
  - データの`.pt`/`.npy`化、CPUメモリ/GPUメモリに乗せておく、とかで若干は速くできるものの…

## 悪魔のささやき『ちゃんぽん』
- そんなときにたまたま`Tensorflow/datasets`の`DataLoader`に出会い、しかもそれが`tf`系の形式じゃなくて`numpy`形式でもイテレートしてくれるそうということを知って魔が差してしまいました…。

# やったこと
- **Pytorch製のモデルと学習器 + Tensorflow製のデータローダー** で、簡単なディープの実験を回してみて、純正Pytorchコードと比較しました。
  - 注意！）条件はだいたい揃えていますが、必ずしも厳密に比較できているわけではない点をご了承ください。
- Tensorflow版のデータローダーには`tf.data.Dataset`系を採用してます。
  - 興味のある方は↓を参考にしてみてください
  - [ドキュメント](https://www.tensorflow.org/api_docs/python/tf/data/Dataset), [GitHub](https://github.com/tensorflow/datasets), [デフォルトで用意されてるデータセット一覧(120種類くらいあります、すごい)](https://www.tensorflow.org/datasets/catalog/overview)



## 環境
手元のzbox(4CPU)で実験しました。

- CUDA: 9.2
- torch:                   1.3.0+cu92
- torchvision:             0.4.2+cu92
- tensorflow:              2.0.0
- tensorflow-datasets      1.3.2

## インストール

```shell
pip install tensorflow-datasets
```
あとは普通に

## 実験
MNIST分類のコードを2種類書いて比較しました。
(参考: [Pytorch本家のサンプルコード](https://github.com/pytorch/examples/blob/master/mnist/main.py)(をだいぶシンプルにして使いました))

- まずは共通部分のPytorchモデルのコード

```python
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.fc1 = nn.Linear(12*12*64, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.max_pool2d(x, 2)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output
```

- 続いて純正Pytorch用のメイン

```python
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
    print('Epoch {} Train loss: {:.6f}; '.format(epoch, loss.item()), end="")


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
    test_loss /= len(test_loader.dataset)
    print('Test loss: {:.4f}, acc: {}/{} ({:.0f}%)'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

### ----------------
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
### ----------------

import time

def main():
    device = torch.device("cuda")
    
    ### ----------------
    transform=transforms.Compose([transforms.ToTensor()])
    kwargs = {'num_workers': 4, 'pin_memory': True}

    train_loader = DataLoader(
        datasets.MNIST('../data', train=True, download=True, transform=transform),
        batch_size=64, shuffle=True, **kwargs)    
    test_loader = DataLoader(
        datasets.MNIST('../data', train=False, download=True, transform=transform),
        batch_size=1000, shuffle=True, **kwargs)
    ### ----------------
    
    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=1.0)

    start = time.time()
    for epoch in range(1, 11):
        train(model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
    elapsed_time = time.time() - start
    print ("elapsed_time:{0}".format(elapsed_time) + "[sec]")

if __name__ == '__main__':
    main()
```

- `### ----------------`で囲まれたところが大きく変わるところです

- 最後に主役のPytorch+Tensorflowのメイン

```python
def train(model, device, train_loader, optimizer, epoch):
    model.train()
    n = 0 ###
    for data, target in train_loader:
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        n += len(data) ###
        if n > 60000: ###
            break ###
    print('Epoch {} Train loss: {:.6f}; '.format(epoch, loss.item()), end="")


def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        n = 0
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            n += len(data) ###
            if n > 10000: ###
                break ###
    test_loss /= n
    print('Test loss: {:.4f}, acc: {}/{} ({:.0f}%)'.format(
        test_loss, correct, n,
        100. * correct / n))

### ----------------
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
tf.compat.v1.enable_eager_execution()

class MnistDataLoader():
    def __init__(self, split, batch_size, epochs):
        self.ds, self.info = tfds.load(name="mnist", split=split, in_memory=True, with_info=True)
        self.num_examples = self.info.splits[split].num_examples
        self.ds = self.ds.repeat(epochs).shuffle(self.num_examples).batch(batch_size)
        self.ds = self.ds.map(
            lambda x:{'image': tf.dtypes.cast(tf.transpose(x['image'], [0,3,1,2]), tf.float32) / 255., 
                      'label': x['label']},
            num_parallel_calls=4)
        self.ds = self.ds.prefetch(5)
        self.ds = tfds.as_numpy(self.ds)
    def __iter__(self):
        return self
    def __next__(self):
        batch = next(self.ds)
        data, target = batch['image'], batch['label']
        return torch.from_numpy(data), torch.from_numpy(target)
    def __len__(self):
        return self.num_examples
### ----------------

import time

def main():
    device = torch.device("cuda")

    ### ----------------
    train_loader = MnistDataLoader(split="train", batch_size=64, epochs=10)
    test_loader = MnistDataLoader(split="test", batch_size=1024, epochs=10)
    ### ----------------
    
    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=1.0)

    start = time.time()
    for epoch in range(1, 11):
        train(model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)
    elapsed_time = time.time() - start
    print ("elapsed_time:{0}".format(elapsed_time) + "[sec]")

if __name__ == '__main__':
    main()
```

- tfds形式のDatasetを軸に、Pytorch風のDataLoaderクラスを定義しています。
- 今回はデフォルトで用意されているMNISTクラスをロードしています。
  - あとに書くように、自分でtfds形式のデータセットを用意することもできます
- `tfds.as_numpy`で、`tensorflow`形式ではなく`np`形式でイテレートできるようになります、神。

## 実行！

（なんと、ちゃんぽんコードは無事に動いてくれました！！！笑）

- 純正Pytorch
  - 70.92678928375244\[sec\](num_workers=1)
  - 43.6055371761322\[sec\](num_workers=4)
  - 48.38671255111694\[sec\](num_workers=8)
  - 54.59042549133301\[sec\](num_workers=16)

- Pytorch + Tensorflow (Ours)
  - **38.35309600830078[sec]**

はい、やはりOutperformeしてくれました:relaxed:
他のサーバーで回してもちゃんぽんが圧勝でした

- Pytorch + Tensorflow
![Screenshot from 2019-12-07 20-23-51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/188173/88fb6e92-d618-7159-399d-2cb603d1be83.png)

- 純正Pytorch
![Screenshot from 2019-12-07 20-25-37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/188173/746ae6d3-9ae8-a588-7e2a-ce1d0a9a355d.png)


## どこで差がついたのか…?
- 実はtfds版のほうが、epochが始まる前(タイマー圏外)にかかる時間は5秒くらい長かったです。
  - がそれはepochを重ねれば誤差の範囲
- tfdsの`.prefetch()`が強そう(しかも簡単！)
  - pytorchにもあるかもしれないけどggってもぱっとしない
- tfdsの`.map()`が強そう(しかも簡単)
  - `__init__`で`.map()`って言いながら実際に計算するのはprefetchの時っぽくて賢そう。

## tfds版はもっと速くできるかも…?
- `ds.cache()`でデータセットのキャッシュ化
- `tf.data.experimental.AUTOTUNE`を使うと、`ds.prefetch()`や`.map()`時の並列化数を自動で決めてくれるらしい。強ひ。
- 若干チートではあるけど、`ds.shuffle()`するときのバッファーサイズ（何個ずつの範囲でシャッフルするか？）を小さくすると、シャッフルが雑になるけど結構速くなる。
  - 少々雑にしても精度は変わらないと思う

## tensorflow-datasetって自作データでも使えるの？

- Yes!
- つい先日、tensorflow-datasetの使い方について素晴らしい記事が出ていたので紹介
  - https://qiita.com/Suguru_Toyohara/items/820b0dad955ecd91c7f3
- ちょっとだけかじらせていただくと、tfds形式のデータローダーは↓な感じで簡単に用意できます

```python
import tensorflow as tf
import tensorflow_datasets as tfds

dataset = tf.data.Dataset.from_tensor_slices((tf.range(10),tf.range(10,20)))
dataset = tfds.as_numpy(dataset)
for item1,item2 in dataset:
    print("data:", item1," target:", item2)
```
![Screenshot from 2019-12-07 20-54-16.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/188173/a906f07c-facc-f106-a5e1-0059ed314abf.png)

- メモリに載り切らないデータを使う場合は、[ここらへん](https://www.tensorflow.org/datasets/add_dataset)を参考に↓を埋めたらtfdsで扱えそうです。

```python
import tensorflow_datasets.public_api as tfds

class MyDataset(tfds.core.GeneratorBasedBuilder):
  """Short description of my dataset."""
  VERSION = tfds.core.Version('0.1.0')
  def _info(self):
    # Specifies the tfds.core.DatasetInfo object
    pass # TODO
  def _split_generators(self, dl_manager):
    # Downloads the data and defines the splits
    # dl_manager is a tfds.download.DownloadManager that can be used to
    # download and extract URLs
    pass  # TODO
  def _generate_examples(self):
    # Yields examples from the dataset
    yield 'key', {}
```

# まとめ
- Pytorchユーザーも、データローダーにはTensorflowを使うというのはアリかもしれないです。たぶん速いです。
- tensorflow-datasetsはデフォルトで(mnistとか以外にも)120種類くらいのデータセットが用意されていて、しかも`np`形式でイテレートしてくれるので使い得な気がします。

# 以上！！

（[twitter](https://twitter.com/StaPriEG2)もよろしく）


```math
\def\textlarge#1{%
  {\rm\Large #1}
}
\def\textsmall#1{%
  {\rm\scriptsize #1}
}
```

　　$\textsmall{い\ い\ ね\ ほ\ し\ い\ ！}$