Pytorch tips

Pytorchで割とよく検索するけど、ひとつの記事にするほどでもない事をまとめてみる。使っていて便利そうだと思った情報はどんどん追加していくつもり。

cuda.is_available(): GPU動作チェック

PytorchでGPUの動作をチェックするにはtorch.cuda.is_available()の値をチェックするだけ。 TrueであればGPUは動作している。

import torch
print(torch.cuda.is_available())

GPUかCPUかを自動判断

以下のおまじないdevice = ’cuda' ...をプログラムのはじめに一行入れておく。そしてモデルと変数を定義したところでto(device)を書いておく。GPU搭載のPCでCPUで動かしたい場合は代わりにdevice = 'cpu'と書けばよい。

import torch
from torchvision import models

device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = models.resnet50().to(device)
x = x.to(device)
y = y.to(device)

to(x.deice): model中での変数がGPUかCPUか

modelの中で新しく変数を定義した場合は、deviceがGPUなのかCPUなのかを指定する必要がある。基本的には自動でどちらでも動くようにすべきであるが、if文などで分ける必要はない。入力されたtensorのdeviceを参照するだけでよい。

def forward(self, x):
    a = torch.randn(x.shape).to(x.device)
    return x*a

numpy ndarray と pytorch tensorの相互変換

こちらの記事にまとめました。 tzmi.hatenablog.com

平均と標準偏差で正規化

Pytorchで使える学習済みモデルへ入力する際の画素値の変換。transform.normarizeを使うのもいいけど、以下のように書けば直接かける。

# x.shape = [n, 3, h, w], range: (0, 1)
mean = torch.tensor([0.485, 0.456, 0.406]).reshape(1,3,1,1)
std = torch.tensor([0.229, 0.224, 0.225]).reshape(1,3,1,1)
eps = 1e-7
x = (x-mean)/(std+eps)

permute: 次元の入れ替え

numpyであれば、transposeを使えば次元の位置をどんどん入れ替えられるが、pytorchでのtranspose()関数は指定した２つの次元の位置を入れ替える関数なので、numpyのtransposeとは挙動が異なる。pytorchでnumpy.transposeと同じことをしたいときはpermute()関数を使う。channel lastからchannel firstへの変換などで使う。

x = torch.randn((4,3,8,8)) # x.shape = [4,3,8,8]
a = x.permute(0,2,3,1) # a.shape= [4,8,8,3]

画像左右反転

pytorchではnumpyのx= x[:,:,::-1]が使えないので、反転した整数配列を使って画像を反転させる。具体的には以下のようにする。

n,c,h,w = x.shape
a = torch.arange(w-1,-1,-1) # start=w-1, end=-1, step=-1
x = x[:,:,:,a]

like関数

numpy のzeors_likeなどと似たことができる関数。torch tensorからshapeやtypeだけでなく、deviceも継承できる。

empty_like, full_like, ones_like, rand_like, randint_like, randn_like, zeros_like

パラメータをもたないモジュール

Swishのようにパラメータを持たないモジュールは以下のように__init__関数を省略できる

class Swish(nn.Module):
    def forward(self, x):
        return x*torch.sigmoid(x)

register_buffer(): 最適化しないパラメータ

pytorchで定義するモジュールの中で、最適化されるパラメータ以外のパラメータを定義することができる。例えば、以下のように書くとパラメータ aは最適化され、パラメータbは最適化されないというように分けることができる。

例えばbatchnorm2dのrunning mean やrunning stdはこのような記法で書かれている。オリジナルのモジュールを作りたいときに重宝するかもしれない。model.parametersでは出てこない(最適化パラメータではない)のに、model.state_dictでは出てくる（保存が可能）という点で、非常に使い勝手がいい。

class A(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Parameter(torch.randn(10))
        self.register_buffer('b', torch.randn(10))

    def forward(self, x):
        return self.a*x + b

model.apply(): モジュールを再帰的に処理

Pytorchのmodel. apply()関数はモデルの中にある全てのモジュールを再帰的に掘って関数を適用できる。例えば、pytorchのresnet50は構造的にresnet50のモジュールの下に全てのconv2dがあるわけではないが、下記のようにすれば全てのconv2dを再帰的に初期化できる。

def weights_init(m):
    if isinstance(m, nn.Conv2d):
        nn.init.kaiming_normal_(m.weight)

model = models.resnet50()
model.apply(weights_init)

tensor.item(): スカラーの取り出し

tensor. item()関数は、tensorの中の変数が1つであることが必要であるが、deviceがGPUでもCPUでも関係なくpython scalarの数値を返してくれる

in [1]: import torch

In [2]: x = torch.tensor([10.0])

In [3]: x
Out[3]: tensor([10.])

In [4]: x.item()
Out[4]: 10.0

model.parameters()の中身のみかた

model.parametersはpythonのgeneratorなので、直接の中を見るときにforループを書く必要がありそうだが、ipythonで以下のようにlist()でラップすれば中身を見れる。

params = list(model.parameters())