はじめに

前回に引き続き、PyTorch 公式チュートリアルの第5弾です。
今回は Learning PyTorch with Examples を進めます。

Learning PyTorch with Examples

このチュートリアルでは、サンプルコードを通じてPyTorchの2つの主な機能を紹介します。

Tensor
自動微分とニューラルネットワーク

サンプルコードで扱うネットワーク（モデル）は3層（入力層、隠れ層 × 1、出力層）です。
活性化関数は ReLU を使用します。

1. Tensor

1.1. Warm-up: numpy

PyTorch の前に、まずは numpy を使用してネットワークを実装します。
Numpy には、ディープラーニング、勾配についての機能はありませんが、
手動で実装することで、簡単なニューラルネットワークを構築できます。

import numpy as np

# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# ランダムな入力データと教師データを作成します
x = np.random.randn(N, D_in)
y = np.random.randn(N, D_out)

# ランダムな値で重みを初期化します
w1 = np.random.randn(D_in, H)
w2 = np.random.randn(H, D_out)

learning_rate = 1e-6
for t in range(500):
    # 順伝播: 現在の重みの値で、予測値 y を計算します
    h = x.dot(w1)
    h_relu = np.maximum(h, 0)
    y_pred = h_relu.dot(w2)

    # ロス（損失）を計算し出力します
    loss = np.square(y_pred - y).sum()
    print(t, loss)

    # ロス値を参考に、逆伝播で 重み w1、w2 の勾配を計算します。
    grad_y_pred = 2.0 * (y_pred - y)
    grad_w2 = h_relu.T.dot(grad_y_pred)
    grad_h_relu = grad_y_pred.dot(w2.T)
    grad_h = grad_h_relu.copy()
    grad_h[h < 0] = 0
    grad_w1 = x.T.dot(grad_h)

    # 重みを更新します。
    w1 -= learning_rate * grad_w1
    w2 -= learning_rate * grad_w2

このコードを実行すると、ロス値が減少し、学習が進んでいることが確認できます。

1.2. PyTorch: Tensors

Numpy はGPUを利用して計算することができませんが、PyTorch の Tensor は GPU を利用して数値計算を高速化できます。
Tensor は勾配も計算できますが、ここでは、上記の numpy の例のように、手動で実装してみます。

import torch


dtype = torch.float
device = torch.device("cpu")
# device = torch.device("cuda:0") # GPU で実行するにはここのコメントを解除します。

# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# ランダムな入力データと教師データを作成します
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# ランダムな値で重みを初期化します
w1 = torch.randn(D_in, H, device=device, dtype=dtype)
w2 = torch.randn(H, D_out, device=device, dtype=dtype)

learning_rate = 1e-6
for t in range(500):
    # 順伝播: 現在の重みの値で、予測値 y を計算します
    h = x.mm(w1)
    h_relu = h.clamp(min=0)
    y_pred = h_relu.mm(w2)

    # ロス（損失）を計算し出力します
    loss = (y_pred - y).pow(2).sum().item()
    if t % 100 == 99:
        print(t, loss)

    # ロス値を参考に、逆伝播で 重み w1、w2 の勾配を計算します。
    grad_y_pred = 2.0 * (y_pred - y)
    grad_w2 = h_relu.t().mm(grad_y_pred)
    grad_h_relu = grad_y_pred.mm(w2.t())
    grad_h = grad_h_relu.clone()
    grad_h[h < 0] = 0
    grad_w1 = x.t().mm(grad_h)

    # 勾配降下法を使用して重みを更新します
    w1 -= learning_rate * grad_w1
    w2 -= learning_rate * grad_w2

このコードでも、ロス値が減少し、学習が進んでいることが確認できます。

2. Autograd

2.1. PyTorch: Tensors and autograd

上記の例では、順伝播と逆伝播を手動で実装しましたが、PyTorch の autograd パッケージを利用すると、逆伝播の計算を自動化できます。

・勾配を計算したい変数(Tensor)の requires_grad = True にする
・backward() を実行する
この2つで逆伝播の計算を自動化できます。

import torch

dtype = torch.float
device = torch.device("cpu")
# device = torch.device("cuda:0") # GPU で実行するにはここのコメントを解除します。

# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# 入力データと教師データを保持するランダムな Tensor を作成します。
# require_grad = False を設定すると、勾配を計算する必要がないことを示します。
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# 重みを保持するランダムな Tensor を作成します。
# requires_grad = True を設定すると、勾配を計算することを示します。
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

learning_rate = 1e-6
for t in range(500):
    # 順伝播: Tensor の演算を利用して 予測値 y を計算します
    # 逆伝播を手動で計算しないため、中間値 h_relu は保持する必要はありません
    y_pred = x.mm(w1).clamp(min=0).mm(w2)

    # Tensorの演算を使用して損失を計算、表示します
    # 損失は形状（1,）のTensorです
    # loss.item() は損失に保持されているスカラー値を取得します
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        print(t, loss.item())

    # autogradを使用して、逆伝播を計算します
    # backward() は、requires_grad = True のすべての Tensor に関する loss の勾配を計算します
    # この呼び出しの後、w1.grad および w2.grad は、それぞれ w1 , w2 の勾配を保持する Tensor になります
    loss.backward()

    # 最急降下法を使用して重みを手動で更新します
    # 重みには require_grad = True があるため、torch.no_grad() で計算グラフが更新されないようにします
    # torch.optim.SGD を使用してこの処理と同じことができます
    with torch.no_grad():
        w1 -= learning_rate * w1.grad
        w2 -= learning_rate * w2.grad

        # 重みを更新した後、手動で勾配をゼロにします
        w1.grad.zero_()
        w2.grad.zero_()

チュートリアルにはありませんが、逆伝播の計算グラフを図示してみます。
torchviz を利用することで計算グラフを図示できます。
colaboratory を使用している場合、インストールが必要です。

!pip install torchviz

PyTorch: Tensors のサンプルコードに少し手を加えます。
ループをやめ、勾配が一度だけ計算されるようにします。

# ランダムな入力データと教師データを作成します
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# ランダムな値で重みを初期化します
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

# 順伝播: 現在の重みの値で、予測値 y を計算します
h = x.mm(w1)
h_relu = h.clamp(min=0)
y_pred = h_relu.mm(w2)

# ロス（損失）を計算し出力します
loss = (y_pred - y).pow(2).sum().item()

# ロス値を参考に、逆伝播で 重み w1、w2 の勾配を計算します。
grad_y_pred = 2.0 * (y_pred - y)
grad_w2 = h_relu.t().mm(grad_y_pred)
grad_h_relu = grad_y_pred.mm(w2.t())
grad_h = grad_h_relu.clone()
grad_h[h < 0] = 0
grad_w1 = x.t().mm(grad_h)

torchviz の make_dot で計算グラフを図にします。
順伝播と勾配を図示します。
param_dict は必須ではありませんが、指定すると図に変数名を記述できます。

# 順伝播の計算グラフを図示します。
from torchviz import make_dot
param_dict = {'w1': w1, 'w2': w2}
make_dot(loss, param_dict)

# w1の勾配の計算グラフを図示します。
make_dot(grad_w1, param_dict)

# w2の勾配の計算グラフを図示します。
make_dot(grad_w2, param_dict)

計算グラフは以下です。

同じように、PyTorch: Tensors and autograd のサンプルコードに手を加え、勾配を一度だけ計算されるようにします。
backward() 実行時に create_graph=True を指定することで導関数のグラフが保持されます。

import torch

# 入力データと教師データを保持するランダムな Tensor を作成します。
# require_grad = False を設定すると、勾配を計算する必要がないことを示します。
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# 重みを保持するランダムな Tensor を作成します。
# requires_grad = True を設定すると、勾配を計算することを示します。
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

# 順伝播: Tensor の演算を利用して 予測値 y を計算します
# 逆伝播を手動で計算しないため、中間値 h_relu は保持する必要はありません
y_pred = x.mm(w1).clamp(min=0).mm(w2)

# Tensorの演算を使用して損失を計算、表示します
# 損失は形状（1,）のTensorです
# loss.item() は損失に保持されているスカラー値を取得します
loss = (y_pred - y).pow(2).sum()

# autogradを使用して、逆伝播を計算します
# backward() は、requires_grad = True のすべての Tensor に関する loss の勾配を計算します
# この呼び出しの後、w1.grad および w2.grad は、それぞれ w1 , w2 の勾配を保持する Tensor になります
loss_backward = loss.backward(create_graph=True)

同じように順伝播と autograd で計算された勾配を図示します。

# 順伝播の計算グラフを図示します。
param_dict = {'w1': w1, 'w2': w2}
make_dot(loss, param_dict)

# w1の勾配の計算グラフを図示します。
make_dot(w1.grad, param_dict)

# w2の勾配の計算グラフを図示します。
make_dot(w2.grad, param_dict)

順伝播は同じです。逆伝播は少し形が異なりますが、autograd で逆伝播の計算が自動で行われていることが確認できます。

2.2. PyTorch: Defining new autograd functions

PyTorchでは、torch.autograd.Functionのサブクラスを定義することで、独自の関数（演算子）を定義できます。
サブクラスには、以下の2つメソッドを実装します。

forward メソッド：入力 Tensor から出力 Tensor を計算する
backward メソッド：出力 Tensor の勾配を受け取り、入力 Tensor の勾配を計算する

この例では、ReLU関数を意味する独自の関数を用いて、2層ネットワークを定義します。

import torch

class MyReLU(torch.autograd.Function):
    """
    torch.autograd.Functionをサブクラス化し、
    Tensorsで動作するフォワードパスとバックワードパスを実装することで、
    独自のカスタム autograd 関数を実装できます。
    """

    @staticmethod
    def forward(ctx, input):
        """
        フォワードパスでは、入力を含むTensorを受け取り、
        出力を含むTensorを返します。 
        ctxは、逆伝播計算のためのオブジェクトです。 
        ctx.save_for_backwardメソッドを使用して、
        オブジェクトをキャッシュできます。
        """
        ctx.save_for_backward(input)
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        """
        backward では、出力に関する損失の勾配を含むTensorを受け取り、
        入力に関する損失の勾配を計算する必要があります。
        """
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input


dtype = torch.float
device = torch.device("cpu")
# device = torch.device("cuda:0") # GPU で実行するにはここのコメントを解除します。

# N     ： バッチサイズ 
# D_in  ： 入力次元数
# H     ： 隠れ層の次元数
# D_out ： 出力次元数
N, D_in, H, D_out = 64, 1000, 100, 10

# 入力データと教師データを保持するランダムな Tensor を作成します。
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# 重みを保持するランダムな Tensor を作成します。
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

learning_rate = 1e-6
for t in range(500):
    # 関数を適用するには、Function.applyメソッドを使用します。
    relu = MyReLU.apply

    # 順伝播: カスタム autograd 関数を利用して 予測値 y を計算します
    y_pred = relu(x.mm(w1)).mm(w2)

    # 損失を計算、表示します
    loss = (y_pred - y).pow(2).sum()
    if t % 100 == 99:
        print(t, loss.item())

    # autogradを使用して、逆伝播を計算します
    loss.backward()

    # 最急降下法を使用して重みを更新します
    with torch.no_grad():
        w1 -= learning_rate * w1.grad
        w2 -= learning_rate * w2.grad

        # 重みを更新した後、手動で勾配をゼロにします
        w1.grad.zero_()
        w2.grad.zero_()

独自関数も可視化してみます。
先ほどと同様、1回のみ処理されるようにします。

# 入力データと教師データを保持するランダムな Tensor を作成します。
x = torch.randn(N, D_in, device=device, dtype=dtype)
y = torch.randn(N, D_out, device=device, dtype=dtype)

# 重みを保持するランダムな Tensor を作成します。
w1 = torch.randn(D_in, H, device=device, dtype=dtype, requires_grad=True)
w2 = torch.randn(H, D_out, device=device, dtype=dtype, requires_grad=True)

# 関数を適用するには、Function.applyメソッドを使用します。
relu = MyReLU.apply

# 順伝播: カスタム autograd 関数を利用して 予測値 y を計算します
y_pred = relu(x.mm(w1)).mm(w2)

# 損失を計算、表示します
loss = (y_pred - y).pow(2).sum()

# autogradを使用して、逆伝播を計算します
loss.backward(create_graph=True)

だいたい、似たような計算グラフになっているでしょうか。

#続く

長くなりましたので、PyTorch: nn は後編に分けたいと思います。

#履歴
2020/05/27 初版公開
2020/07/10 後編のリンク追加