MLP-Mixer

MLP-Mixer: An all-MLP Architecture for Vision

はじめに

　近年、自然言語処理・画像認識・音声認識において、Transformerを使った事前学習モデルの成果が著しいが、一方で古典的なMLP(Multi Layer Perceptron)を改良した驚くほどシンプルなアーキテクチャでTransformerと同等の性能がでることがMLP-Mixer: An all-MLP Architecture for Visionで確認されており、非常に気になるところである。本記事では、MLP-MixerとVision Transformerの推論の精度をImagenetのデータセットで確認してみる。

対象読者

人工知能・機械学習・深層学習の概要は知っているという方
理論より実装を重視する方
Pythonを触ったことがある方
Pytorchを触ったことがある方

1. 論文サマリ

MLP-Mixer: An all-MLP Architecture for Visionが主張しているポイントを
三つに絞るとすると、以下の通りだ。

CNNやAttentionは画像認識のベンチマークで高精度な記録を出すのに重要だが必須ではない。
２つの”まぜまぜ”層によってMLPだけでもSoTAに匹敵する性能を出せる。
- 画像の位置情報（画像パッチ）のまぜまぜ。
- 画像の空間的情報（画像パッチまたぎの情報）のまぜまぜ。
訓練と推論にかかるコストもSoTAに匹敵する性能をだせる。

詳細は論文に、概要はAI-Scholarさんの記事"【MLP-Mixer】MLPがCNN，Transformerを超える日"におまかせ。本記事では実際にGoogle Colaboratoryで使う方法を解説する。

結果比較

MLP-MixerとViTどっちがすごいの？って結果だけ知りたいビジネスチックな方のために
結論ファースト。cipher10のテストにおいて、事前学習モデルでそのまま
比較すると全く同精度。再学習モデルでの比較ではViTがちょっとだけ（1%くらい）まだ高精度。

パラメータサイズはMLP-Mixerのが3/4くらいで省エネなので、自分はMLP-Miexer推し。

この研究をきっかけにgMLPが"Pay attention to MLPs"という
Attention意識した論文ではattentionも組み合わせて更に精度が上がっている模様。Attention is Not All You Needというわけだ。

Model	事前学習モデル正解率(Cipher10)	再学習後モデル正解率(Cipher10)
Vision Transformer(ViT)	0.10063	0.976
MLP-Mixer	0.10063	0.9646382

2. 画像認識で比較

まずは、Vision Transformer(ViT)で画像認識。

セットアップ

ColabランタイムはGPUにする。

リポジトリをvision_transformerというディレクトリにクローンして最新化

![ -d vision_transformer ] || git clone --depth=1 https://github.com/google-research/vision_transformer
!cd vision_transformer && git pull

!pip install -qr vision_transformer/vit_jax/requirements.txt```

利用可能な事前学習モデルを表示

!gsutil ls -lh gs://vit_models/imagenet*
!gsutil ls -lh gs://vit_models/sam
!gsutil ls -lh gs://mixer_models/*

gs://vit_models/imagenet21k+imagenet2012/:
377.57 MiB  2020-11-30T16:17:02Z  gs://vit_models/imagenet21k+imagenet2012/R50+ViT-B_16.npz
330.29 MiB  2020-10-29T17:05:52Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_16-224.npz
 331.4 MiB  2020-10-20T11:48:22Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_16.npz
336.89 MiB  2020-10-20T11:47:36Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_32.npz
334.78 MiB  2021-03-12T09:04:16Z  gs://vit_models/imagenet21k+imagenet2012/ViT-B_8.npz
  1.13 GiB  2020-10-29T17:08:31Z  gs://vit_models/imagenet21k+imagenet2012/ViT-L_16-224.npz
  1.14 GiB  2020-10-20T11:53:44Z  gs://vit_models/imagenet21k+imagenet2012/ViT-L_16.npz
  1.14 GiB  2020-10-20T11:50:56Z  gs://vit_models/imagenet21k+imagenet2012/ViT-L_32.npz

gs://vit_models/imagenet21k/:
450.23 MiB  2021-01-20T14:12:43Z  gs://vit_models/imagenet21k/R26+ViT-B_32.npz
439.85 MiB  2020-11-30T10:10:15Z  gs://vit_models/imagenet21k/R50+ViT-B_16.npz
  1.31 GiB  2021-01-20T14:11:54Z  gs://vit_models/imagenet21k/R50+ViT-L_32.npz
393.69 MiB  2020-10-22T21:38:39Z  gs://vit_models/imagenet21k/ViT-B_16.npz
400.01 MiB  2020-11-02T08:30:56Z  gs://vit_models/imagenet21k/ViT-B_32.npz
393.72 MiB  2021-03-10T13:28:28Z  gs://vit_models/imagenet21k/ViT-B_8.npz
  2.46 GiB  2020-11-03T10:46:11Z  gs://vit_models/imagenet21k/ViT-H_14.npz
  1.22 GiB  2020-11-09T14:39:51Z  gs://vit_models/imagenet21k/ViT-L_16.npz
  1.23 GiB  2020-11-02T08:35:10Z  gs://vit_models/imagenet21k/ViT-L_32.npz
TOTAL: 17 objects, 14306096550 bytes (13.32 GiB)
 330.3 MiB  2021-07-13T19:39:09Z  gs://vit_models/sam/ViT-B_16.npz
336.61 MiB  2021-07-13T19:39:10Z  gs://vit_models/sam/ViT-B_32.npz
  1.13 GiB  2021-07-13T19:39:38Z  gs://vit_models/sam/ViT-L_16.npz
  1.14 GiB  2021-07-13T19:39:38Z  gs://vit_models/sam/ViT-L_32.npz
TOTAL: 4 objects, 3143025464 bytes (2.93 GiB)
       6 B  2021-06-28T13:07:12Z  gs://mixer_models/sam_$folder$

gs://mixer_models/imagenet1k/:
228.47 MiB  2021-05-05T14:09:01Z  gs://mixer_models/imagenet1k/Mixer-B_16.npz
794.29 MiB  2021-05-05T14:09:02Z  gs://mixer_models/imagenet1k/Mixer-L_16.npz

gs://mixer_models/imagenet21k/:
289.61 MiB  2021-05-05T14:09:11Z  gs://mixer_models/imagenet21k/Mixer-B_16.npz
875.78 MiB  2021-05-05T14:09:12Z  gs://mixer_models/imagenet21k/Mixer-L_16.npz

gs://mixer_models/sam/:
228.47 MiB  2021-06-28T13:08:09Z  gs://mixer_models/sam/Mixer-B_16.npz
230.04 MiB  2021-06-28T13:08:08Z  gs://mixer_models/sam/Mixer-B_32.npz
TOTAL: 7 objects, 2775222110 bytes (2.58 GiB)

ViTとMixerをダウンロード

model_name = 'ViT-B_32'  #@param ["ViT-B_32", "Mixer-B_16"]

if model_name.startswith('ViT'):
  ![ -e "$model_name".npz ] || gsutil cp gs://vit_models/imagenet21k/"$model_name".npz .
if model_name.startswith('Mixer'):
  ![ -e "$model_name".npz ] || gsutil cp gs://mixer_models/imagenet21k/"$model_name".npz .

Colab上でモデルをプルダウンで選択できて便利。まずは、ViTでいってみよう。

TPU利用できるか確認（→今のところ無料プランではできない）。GPUでいこう。

import os
if 'google.colab' in str(get_ipython()) and 'COLAB_TPU_ADDR' in os.environ:
  import jax
  import jax.tools.colab_tpu
  jax.tools.colab_tpu.setup_tpu()
  print('Connected to TPU.')
else:
  print('No TPU detected. Can be changed under "Runtime/Change runtime type".')

No TPU detected. Can be changed under "Runtime/Change runtime type".

ログ設定とデバイス確認


from absl import logging
import flax
import jax
from matplotlib import pyplot as plt
import numpy as np
import tqdm

logging.set_verbosity(logging.INFO)

# 利用可能なデバイスの数を表示
jax.local_devices()

INFO:absl:Unable to initialize backend 'tpu_driver': Not found: Unable to find driver in registry given worker:
INFO:absl:Unable to initialize backend 'tpu': Invalid argument: TpuPlatform is not available.
[GpuDevice(id=0, process_index=0)]

GPUは使えるよ。

Colabからソースを確認しながら実行する小技：画面右スプリットエディタでいくつかのコードファイルを開く。


from google.colab import files
files.view('vision_transformer/vit_jax/configs/common.py')
files.view('vision_transformer/vit_jax/configs/models.py')
files.view('vision_transformer/vit_jax/checkpoint.py')
files.view('vision_transformer/vit_jax/input_pipeline.py')
files.view('vision_transformer/vit_jax/models.py')
files.view('vision_transformer/vit_jax/momentum_clip.py')
files.view('vision_transformer/vit_jax/train.py')

モジュールのインポート


import sys
if './vision_transformer' not in sys.path:
  sys.path.append('./vision_transformer')

%load_ext autoreload
%autoreload 2

from vit_jax import checkpoint
from vit_jax import input_pipeline
from vit_jax import utils
from vit_jax import models
from vit_jax import momentum_clip
from vit_jax import train
from vit_jax.configs import common as common_config
from vit_jax.configs import models as models_config

Cipher10/100の画像扱うためのHelper


# Helper functions for images.

labelnames = dict(
  # https://www.cs.toronto.edu/~kriz/cifar.html
  cifar10=('airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'),
  # https://www.cs.toronto.edu/~kriz/cifar.html
  cifar100=('apple', 'aquarium_fish', 'baby', 'bear', 'beaver', 'bed', 'bee', 'beetle', 'bicycle', 'bottle', 'bowl', 'boy', 'bridge', 'bus', 'butterfly', 'camel', 'can', 'castle', 'caterpillar', 'cattle', 'chair', 'chimpanzee', 'clock', 'cloud', 'cockroach', 'couch', 'crab', 'crocodile', 'cup', 'dinosaur', 'dolphin', 'elephant', 'flatfish', 'forest', 'fox', 'girl', 'hamster', 'house', 'kangaroo', 'computer_keyboard', 'lamp', 'lawn_mower', 'leopard', 'lion', 'lizard', 'lobster', 'man', 'maple_tree', 'motorcycle', 'mountain', 'mouse', 'mushroom', 'oak_tree', 'orange', 'orchid', 'otter', 'palm_tree', 'pear', 'pickup_truck', 'pine_tree', 'plain', 'plate', 'poppy', 'porcupine', 'possum', 'rabbit', 'raccoon', 'ray', 'road', 'rocket', 'rose', 'sea', 'seal', 'shark', 'shrew', 'skunk', 'skyscraper', 'snail', 'snake', 'spider', 'squirrel', 'streetcar', 'sunflower', 'sweet_pepper', 'table', 'tank', 'telephone', 'television', 'tiger', 'tractor', 'train', 'trout', 'tulip', 'turtle', 'wardrobe', 'whale', 'willow_tree', 'wolf', 'woman', 'worm')
)
def make_label_getter(dataset):
  """Returns a function converting label indices to names."""
  def getter(label):
    if dataset in labelnames:
      return labelnames[dataset][label]
    return f'label={label}'
  return getter

def show_img(img, ax=None, title=None):
  """Shows a single image."""
  if ax is None:
    ax = plt.gca()
  ax.imshow(img[...])
  ax.set_xticks([])
  ax.set_yticks([])
  if title:
    ax.set_title(title)

def show_img_grid(imgs, titles):
  """Shows a grid of images."""
  n = int(np.ceil(len(imgs)**.5))
  _, axs = plt.subplots(n, n, figsize=(3 * n, 3 * n))
  for i, (img, title) in enumerate(zip(imgs, titles)):
    img = (img + 1) / 2  # Denormalize
    show_img(img, axs[i // n][i % n], title)

データセットのロード


dataset = 'cifar10'
batch_size = 512
config = common_config.with_dataset(common_config.get_config(), dataset)
num_classes = input_pipeline.get_dataset_info(dataset, 'train')['num_classes']
config.batch = batch_size
config.pp.crop = 224

INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: cifar10/3.0.2
INFO:absl:Load dataset info from /tmp/tmpmx7kqa7vtfds
INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.

データセットの設定については、右記のinput_pipeline.pyを参照


ds_train = input_pipeline.get_data_from_tfds(config=config, mode='train')
ds_test = input_pipeline.get_data_from_tfds(config=config, mode='test')

del config  #  データセットをインスタンス化するためにのみ必要です。

INFO:absl:Load pre-computed DatasetInfo (eg: splits, num examples,...) from GCS: cifar10/3.0.2
INFO:absl:Load dataset info from /tmp/tmp5qpln01btfds
INFO:absl:Field info.citation from disk and from code do not match. Keeping the one from code.
INFO:absl:Generating dataset cifar10 (/root/tensorflow_datasets/cifar10/3.0.2)
Downloading and preparing dataset cifar10/3.0.2 (download: 162.17 MiB, generated: 132.40 MiB, total: 294.58 MiB) to /root/tensorflow_datasets/cifar10/3.0.2...
Dl Completed...: 100%
1/1 [00:20<00:00, 18.14s/ url]
Dl Size...: 100%
162/162 [00:20<00:00, 6.62 MiB/s]
Extraction completed...: 100%
1/1 [00:20<00:00, 20.50s/ file]
INFO:absl:Downloading https://www.cs.toronto.edu/~kriz/cifar-10-binary.tar.gz into /root/tensorflow_datasets/downloads/cs.toronto.edu_kriz_cifar-10-binaryODHPtIjLh3oLcXirEISTO7dkzyKjRCuol6lV8Wc6C7s.tar.gz.tmp.60fc73cd82d24890985164371806add3...
INFO:absl:Generating split train

Shuffling and writing examples to /root/tensorflow_datasets/cifar10/3.0.2.incompleteEIKKBP/cifar10-train.tfrecord
100%
49999/50000 [00:00<00:00, 122648.19 examples/s]
INFO:absl:Done writing /root/tensorflow_datasets/cifar10/3.0.2.incompleteEIKKBP/cifar10-train.tfrecord. Shard lengths: [50000]
INFO:absl:Generating split test
Shuffling and writing examples to /root/tensorflow_datasets/cifar10/3.0.2.incompleteEIKKBP/cifar10-test.tfrecord
100%
9999/10000 [00:00<00:00, 46586.49 examples/s]
INFO:absl:Done writing /root/tensorflow_datasets/cifar10/3.0.2.incompleteEIKKBP/cifar10-test.tfrecord. Shard lengths: [10000]
INFO:absl:Skipping computing stats for mode ComputeStatsMode.SKIP.
INFO:absl:Constructing tf.data.Dataset for split train[:98%], from /root/tensorflow_datasets/cifar10/3.0.2
Dataset cifar10 downloaded and prepared to /root/tensorflow_datasets/cifar10/3.0.2. Subsequent calls will reuse this data.
INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
INFO:absl:Reusing dataset cifar10 (/root/tensorflow_datasets/cifar10/3.0.2)
INFO:absl:Constructing tf.data.Dataset for split test, from /root/tensorflow_datasets/cifar10/3.0.2
INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2

説明用にテスト画像を一括して取得


batch = next(iter(ds_test.as_numpy_iterator()))
# 形の意味: [num_local_devices, local_batch_size, h, w, c]
batch['image'].shape

(1, 512, 224, 224, 3)

画像サンプルを表示


images, labels = batch['image'][0][:9], batch['label'][0][:9]
titles = map(make_label_getter(dataset), labels.argmax(axis=1))
show_img_grid(images, titles)

訓練用画像も確認していく。

# 画像がどのように切り取られ、どのように拡大されるかを確認
# 右側のエディタでinput_pipeline.get_data()をチェックすると、 
# 画像の前処理がどのように異なるかがわかる。
batch = next(iter(ds_train.as_numpy_iterator()))
images, labels = batch['image'][0][:9], batch['label'][0][:9]
titles = map(make_label_getter(dataset), labels.argmax(axis=1))
show_img_grid(images, titles)

事前学習モデルロード


model_config = models_config.MODEL_CONFIGS[model_name]
model_config

classifier: token
hidden_size: 768
name: ViT-B_32
patches:
size: !!python/tuple
- 32
- 32
representation_size: null
transformer:
attention_dropout_rate: 0.0
dropout_rate: 0.0
mlp_dim: 3072
num_heads: 12
num_layers: 12

ViTがロードされていて、Transformerが12層あることが確認できるよ。

続いて、パラメータ初期化

# モデルの定義を読み込み、ランダムなパラメータを初期化。
# また、モデルをXLAにコンパイルします（初回は数分かかる）。

if model_name.startswith('Mixer'):
  model = models.MlpMixer(num_classes=num_classes, **model_config)
else:
  model = models.VisionTransformer(num_classes=num_classes, **model_config)
variables = jax.jit(lambda: model.init(
    jax.random.PRNGKey(0),
    # Discard the "num_local_devices" dimension of the batch for initialization.
    batch['image'][0, :1],
    train=False,
), backend='cpu')()

チェックポイントのロード


# チェックポイントのロードと変換。
# 実際に事前学習したモデルの結果をロードしますが、同時に
# 最終層の変更や、位置埋め込みのサイズ変更など、パラメータを少し変更します。
# 位置埋め込みを変更します。
# 詳細については、コードと論文のメソッドを参照のこと。

params = checkpoint.load_pretrained(
    pretrained_path=f'{model_name}.npz',
    init_params=variables['params'],
    model_config=model_config,
)

INFO:absl:Inspect extra keys:
{'pre_logits/bias', 'pre_logits/kernel'}
INFO:absl:load_pretrained: drop-head variant

ここまでで、すべてのデータはホストメモリにある。


# 配列をデバイスに複製する。
# これにより、pytree paramsに含まれるすべての配列がShardedDeviceArrayになる。
# 同じデータがすべてのローカルデバイスに複製される。
# 単一のGPUの場合は、単にデータをデバイスに移動する。

params_repl = flax.jax_utils.replicate(params)
print('params.cls:', type(params['head']['bias']).__name__,
      params['head']['bias'].shape)
print('params_repl.cls:', type(params_repl['head']['bias']).__name__,
      params_repl['head']['bias'].shape)

params.cls: DeviceArray (10,)
params_repl.cls: _ShardedDeviceArray (1, 10)
/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:391: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.
"jax.host_count has been renamed to jax.process_count. This alias "
/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:378: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.
"jax.host_id has been renamed to jax.process_index. This alias "

そして、モデルのフォワードパスの呼び出しを、利用可能なすべてのデバイスにマッピングする。


vit_apply_repl = jax.pmap(lambda params, inputs: model.apply(
    dict(params=params), inputs, train=False))

正解率を計測する用に関数を定義する。


def get_accuracy(params_repl):
  """Returns accuracy evaluated on the test set."""
  good = total = 0
  steps = input_pipeline.get_dataset_info(dataset, 'test')['num_examples'] // batch_size
  for _, batch in zip(tqdm.trange(steps), ds_test.as_numpy_iterator()):
    predicted = vit_apply_repl(params_repl, batch['image'])
    is_same = predicted.argmax(axis=-1) == batch['label'].argmax(axis=-1)
    good += is_same.sum()
    total += len(is_same.flatten())
  return good / total

とりあえず、ViTを事前学習モデルのまま再学習なしにcipher10で正解率計測してみる。


# 再学習していない状態での精度
get_accuracy(params_repl)

INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
100%|██████████| 19/19 [01:25<00:00, 4.49s/it]
DeviceArray(0.10063734, dtype=float32)

10%という結果。そりゃあそうか。再学習してないもの。
ということで再学習してみる。


total_steps = 100
warmup_steps = 5
decay_type = 'cosine'
grad_norm_clip = 1

#accum_steps = 8
accum_steps = 32
base_lr = 0.03

# 詳細は、右側のエディタでtrain.make_update_fnを確認のこと。
lr_fn = utils.create_learning_rate_schedule(total_steps, base_lr, decay_type, warmup_steps)
update_fn_repl = train.make_update_fn(
    apply_fn=model.apply, accum_steps=accum_steps, lr_fn=lr_fn)

# モメンタムで半分の精度を使用しメモリを節約するグラデーションクリッピングを利用。
opt = momentum_clip.Optimizer(grad_norm_clip=grad_norm_clip).create(params)
opt_repl = flax.jax_utils.replicate(opt)

# ドロップアウトのためPRNGsを初期化
update_rng_repl = flax.jax_utils.replicate(jax.random.PRNGKey(0))

/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:391: UserWarning: jax.host_count has been renamed to jax.process_count. This alias will eventually be removed; please update your code.
"jax.host_count has been renamed to jax.process_count. This alias "
/usr/local/lib/python3.7/dist-packages/jax/lib/xla_bridge.py:378: UserWarning: jax.host_id has been renamed to jax.process_index. This alias will eventually be removed; please update your code.
"jax.host_id has been renamed to jax.process_index. This alias "

TPUで20分程度かかるらしいがGPUでも30分くらいだった。


losses = []
lrs = []

for step, batch in zip(
    tqdm.trange(1, total_steps + 1),
    ds_train.as_numpy_iterator(),
):

  opt_repl, loss_repl, update_rng_repl = update_fn_repl(
      opt_repl, flax.jax_utils.replicate(step), batch, update_rng_repl)
  losses.append(loss_repl[0])
  lrs.append(lr_fn(step))

plt.plot(losses)
plt.figure()
plt.plot(lrs)

さて、再学習後の正解率は


# 精度計測
print(model_name)
print(get_accuracy(opt_repl.target))

INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
ViT-B_32
100%|██████████| 19/19 [00:50<00:00, 2.66s/it]
0.9766653

97.6%と結構高い。さすがViT!!!

つづいてMLP-Mixerで画像認識。変えるのはここだけ。

モデルのサイズがまず3/4くらいになったのがうれしい。

まず、事前学習モデルのままの精度はこちら。


# 再学習していない状態での精度
get_accuracy(params_repl)

INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
100%|██████████| 19/19 [02:37<00:00, 8.28s/it]
DeviceArray(0.10063734, dtype=float32)

再学習の損失をプロットするとこんな感じ


losses = []
lrs = []
# GPUで35分程度かかる
for step, batch in zip(
    tqdm.trange(1, total_steps + 1),
    ds_train.as_numpy_iterator(),
):

  opt_repl, loss_repl, update_rng_repl = update_fn_repl(
      opt_repl, flax.jax_utils.replicate(step), batch, update_rng_repl)
  losses.append(loss_repl[0])
  lrs.append(lr_fn(step))

plt.plot(losses)
plt.figure()
plt.plot(lrs)

再学習後の精度はこちら。


# 精度計測
print(model_name)
print(get_accuracy(opt_repl.target))

INFO:absl:Load dataset info from /root/tensorflow_datasets/cifar10/3.0.2
Mixer-B_16
100%|██████████| 19/19 [02:30<00:00, 7.91s/it]
0.9646382

長い記事を読んでいただき感謝です。

参考文献

-MLP-Mixer

-Vision Transformer / MLP-Mixer

著者

ツイッターもやってます。
@keiji_dl

その他オススメ

人工知能Webアプリ開発入門

画像認識をViT(Vision Transformer)を用いて行うWebアプリをPython/Pytorch Lightning Flask/Jinja2/Bootstrap/JQuery/CSS/HTMLなどをフルスタックに使って7ステップ2時間で解説する動画を開設しました。

7ステップで作るPython x Flask x Pytorch 人工知能Webアプリ開発入門

Transformer 101本ノック(Kindle Unlimited)。

ソースコードへのリンクも手に入りますのでぜひご一読下さい。

MLP-MixerのすごさをPython Pytorch実装しながら体感してみる。