【PyTorch】地味に知っておくべき実装の躓きドコロ

はじめに

最近、TensorFlowのEagerが登場したため、Debugの難しさという欠点が緩和される見込みがあります。そこで私自身はPyTorchからTensorFlowに移行しようかなと思っているところです。

s0sem0y.hatenablog.com

その前にPyTorchでところどころ実装時に躓いたりした部分をまとめておきたいと思います。PyTorchに自分自身が戻ってきたいと思った時、あるいはこれからPyTorchを始めるという方の役に立てればと思います。

一応PyTorchで簡単な計算やニューラルネットが書ける程度の知識を有している前提とします。

s0sem0y.hatenablog.com

s0sem0y.hatenablog.com

s0sem0y.hatenablog.com

自動微分に関すること

まず自動微分に関することで抑えておかなければならないことがあります。

PyTorchはいわゆるDefine by Runという仕組みで動作します。簡単に言えば、実際に数値を計算する度に計算の履歴を構築し、微分に必要な情報を保持するという仕組みです。この仕組みの利点は、ニューラルネットの構造（あるいは計算の流れ）がデータごとに毎回異なっていても良いということです。

一方でDefine and Runの場合ニューラルネットの構造（あるいは計算の流れ）を予め構築しておくことで、微分に必要な情報が静的に定まった状態になります。この仕組みでは一度計算を定義したら後はデータを投げるだけで、ほとんど考えることは無くなりますが、データごとに計算を変えることはできません。

Define by Runは利点ばかりが強調されますが、実装を少し間違えるとメモリの使用量が膨れ上がってしまう危険性を持っています。具体的には微分の計算が必要ないような場所で、Variableを用いた計算を大量に行ってしまっている場合などです。

計算が行われる度に、その計算の履歴を構築していくわけですから、不要な履歴のためにメモリをガンガン使っていくことになります。

全バッチのLossを計算

通常はミニバッチ学習を用いるため、バッチ全体に対するLossを算出したければ、ミニバッチのLossを次々に加算していき値を求めることになります。このLossをVariableのままで計算してしまうと、プログラムはこの全ての計算履歴を保持しようとしてしまい、メモリを食ってしまいます。

従ってLossをただ見たいだけならば、一旦VariableからTensorの型だけを取り出して計算する方が効率が良いということになります。例えば以下のような形式になるでしょう。

def train():
    model.train()
    running_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):

        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()

        running_loss += loss.data[0]

running_lossの中身は、TorchTensorの型になっています。loss.data[0]の加算は計算グラフの構築が行われないため無駄に資源を食うことはありません。同様にAccuracyやRecall、Precisionなどを計算する際にも注意が必要です。

loss.data[0]

の0というIndexにアクセスする必要性があるのかは分かりません。

loss.data

のままだと、値とTensorの型の詳細が一緒に表示されます。

loss.data[0]

ですと、値のみが扱われます。基本的に見たいのはLossの値ですから「loss.data[0]」を使うのが定石のようです（これは本当に慣れなかった）。

とにかく微分が不要な計算はVariableではなくTensorで計算をするということです。

Validation時に計算グラフを保持しない

さて、とある計算において全く計算グラフを必要としないケースもあります。

Validationあるいは学習後のTest時には、パラメータの更新をしないのですから微分の計算も不要です。であれば、計算グラフを保持するためにメモリを使う必要もないはずです。

このようなことを指定するためにVariableにはvolatileという引数が用意されています。例えば、以下のようにすることで、dataとtargetにまつわる計算の履歴は保持されなくなります。これで、メモリを節約できますね。

def val():
    model.eval()
    running_loss = 0
    for batch_idx, (data, target) in enumerate(train_loader):

        data = Variable(data, volatile=True)
        target = Variable(target, volatile=True)
        output = model(data)
        loss = F.nll_loss(output, target)
     
        running_loss += loss.data[0]

特にTest時（学習後の評価）においては、volatileを入れたらデータのサイズによってはDataLoaderを使わなくとも、バッチ全体で一気に順伝播が行える場合もありますので試して見てもいいでしょう。

VariableのRank変更

畳み込み層やLSTM、全結合層が入り混じった構造を扱うときには、データのRankを変更しなければなりません。numpyで言うところのreshapeです。

PyTorchではviewというメソッドで変更が可能になります。

## numpy
X = np.random.randn(10, 8)
X = X.reshape(10, 4, 2)

## PyTorch
X = torch.randn(10, 8)
X = X.view(10, 4, 2)  

## PyTorch Variable
X = Variable(torch.randn(10, 8))
X = X.view(10, 4, 2)

Variableはtorch.Tensorと同様にviewメソッドを有しており、ほとんど区別なしに用いることができます。多くの場合、ひとつ目の数字（上記で10）はバッチサイズを表しており、これは固定で使う場合がほとんどであるため、以下のような用い方をするのが便利だと思われます。

## numpy
X = np.random.randn(10, 8)
X = X.reshape(-1, 4, 2)

## PyTorch
X = torch.randn(10, 8)
X = X.view(-1, 4, 2)  

## PyTorch Variable
X = Variable(torch.randn(10, 8))
X = X.view(-1, 4, 2)

viewメソッドが使えないケース（メモリの同一ブロックに格納されていない）

PyTorchを使ってて一度躓いた独特の現象として、Tensor（あるいはVariable）の各値がメモリの同一ブロックに格納されていないというエラーです。

RuntimeError: input is not contiguous

と表示されます。こいつの対処法は、Tensorの各値が同一ブロックに収まった形でコピーすることです。もしもviewメソッドが通らなかった場合には

X = Variable(torch.randn(10, 8))
X = X.contiguous()
X = X.view(-1, 4, 2)

とcontiguous()メソッドを挟みましょう。

contiguous()は、もしも同一ブロックに値が存在しなければ、同一ブロックに収まるようにコピーした上で値を戻します。同一ブロックに既に存在する場合には何もせず値を返します（おいおい、こんくらい自動でやってくれよ…と思ってしまったが）。

出現パターンとしては、畳み込み層でガシガシ計算した後に、全結合層に入力するためにRankを変更する場合に生じることが多かったです。

（contiguous()が破壊的メソッドだったかどうかは忘れました…。破壊的なら代入は要りません）

ドロップアウト関数とドロップアウトクラス

ChainerからPyTorchに移った時に、「んお？？」ってなったのがここです。Chainerでは（今は知りませんが）、関数を使ってドロップアウトを実装するのが一般的でした（クラスもありますが）。Chainerでは主にパラメータを持つものがlinksとして、演算を行うだけのものがfunctionalsとして準備されていたため、ここらへんはChainerが一貫性あり分かりやすい印象でした。

ところがPyTorchはパラメータを持つ層とそうでない層を特に強く区別しません。基本的には全てクラスとして、__init__で準備しておき、fowardでメソッドとして演算を呼び出すというのが定石です。

PyTorchではmodel.train()やmodel.eval()によって、モデルのモードを切り替えますが、これらのメソッドによってドロップアウトを行うか否かを自動で切り替えてくれるのはドロップアウトクラス（torch.nn.Dropout）の方です。torch.nn.functional.dropoutの方は、model.eval()などが働きません。関数の引数でtorch.nn.functional.dropout(training=False)などとしなければならないのです。​_

ちなみにドロップアウト率はデフォルトで0.5となっています。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        ## ドロップアウトクラスのメソッドを利用
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        ## ドロップアウト関数を利用。training引数を必要とする。
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

これはChainerのドロップアウトクラスと関数との関係とちょうど真逆となっています。Chainer触っていた人は注意してください。

使い分け？？

特にこだわりがなければクラスを使うのが安全だと思われます。関数を使う主だった強い理由は今のところ見当たりませんが、例として以下のようなケースが考えられます。

順伝播毎に異なるドロップアウト率を採用する

このようなことが学習を上手く動作させることに繋がるかは知りませんが、実装としてはこのような使い方が考えられます。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        ## 0~1の一様乱数
        p = torch.rand(1)
        ## 順伝播毎に異なるドロップアウト率を利用
        x = F.dropout(x, p, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

ドロップアウトを用いたベイズ学習とベイズ予測分布の近似

あとはmodel.eval()でモードを切り替えたとしても、ドロップアウトは利用し続けたい場合には関数の方を使うことになります。このようなケースはドロップアウトをベイズ推論として解釈して使う場合には有効だと思われます。例えば出力1次元の回帰なら以下のような感じになるでしょうか（ごめんなさい、試してないです）。

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(5, 50)
        self.fc2 = nn.Linear(50, 1)
        self.n_samples = 100

    def forward(self, x):
        x = F.relu(self.fc1(x)) 
        ## 常にドロップアウトを利用
        x = F.dropout(x, training=True)
        x = self.fc2(x)
        return x

    def bayes_forward(self, x)
        samples = []
        for _ in range(self.n_samples):
            samples.append(self(x))
        samples = torch.stack(samples)
        return samples.mean(0), samples.var(0)

 

 

自分で層を定義する

パラメータの初期値を与える方法は数多く提案されています。用いるデータのスケールや活性化関数に応じて、初期化を自分で行いたい場合に起こりうるミスを紹介します。

全結合層を作る

簡単のためいわゆる全結合層（線形層）を自分で作ってみましょう。基本的には以下のような形式になります。nn.Parameterにお望みの形のtorch.Tensorを渡してあげて、適宜reset_parametersメソッドを定義してあげればいいでしょう。

これが基本的な層の定義の仕方だと思われます。単純な掛け算や引き算、要素ごとの積など（何に使えるかは知りませんが）などを定義してやることも可能です。

class MyLinear(nn.Module):

    def __init__(self, in_features, out_features):
        super(MyLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Parameter(
            torch.Tensor(out_features, in_features))
        self.bias = Parameter(torch.Tensor(out_features))
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input):
        return F.linear(input, self.weight, self.bias)

層のパラメータの定義はVariableを使って以下のように定義することも可能です。

class MyLinear(nn.Module):

    def __init__(self, in_features, out_features, bias=True):
        super(MyLinear, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = Variable(
            torch.Tensor(out_features, in_features),
            requires_grad=True)
        self.bias = Variable(
            torch.Tensor(out_features), 
            requires_grad=True)
        self.reset_parameters()

    def reset_parameters(self):
        stdv = 1. / math.sqrt(self.weight.size(1))
        self.weight.data.uniform_(-stdv, stdv)
        self.bias.data.uniform_(-stdv, stdv)

    def forward(self, input):
        return F.linear(input, self.weight, self.bias)

ここで、Variableの引数にrequires_grad=Trueなる引数を与えていることに注意してください。この引数を与えることで、パラメータの更新の際に勾配を使うということを指定することができます。

仮にこの引数を与えなかった場合、勾配はNoneとなり、パラメータの更新が行われません。Variableの方は入力データ、出力データ、LSTMの内部状態など、PyTorch内で広く一般的に使われるデータの型です。Variableの型が与えられているものの多くはパラメータではないため、学習による更新は行われません。従ってVariableのrequires_gradはデフォルトでFalseとなっています。

実際にはパラメータを自分で決めて使うときには

「from torch.nn import Parameter」としてParameterを使うと決めておけば、このような困ったことも起こりづらいはずです（この点、TensorFlowには更新が行われるテンソルに対してはtf.Variableで包むという決まりですので分かりやすいです）。

最後に

なんか他にもデータの型で怒られたりいろいろあった気がしますが、結構出くわす機会が多いので、すぐに解決方法が見つかると思います。今回は、あんまり意識しないタイプの知っておくべきことを抑えておきました。

次はTensorFlowのEagerモードにボコボコにされたいと思います。