Pythonいぬ

pythonを使った画像処理に関する記事を書いていきます

Pytorchでtensorのリサイズ

pytorch で tensor の画像サイズをリサイズするとき、numpyなどに変換して画像リサイズしてから逆変換することがよくある。しかし、学習の途中でリサイズする場合は numpyに戻さずにリサイズしないといけない。こういう場合は、F.interpolatenn.Upsample, F.adaptive_ave_pool2dを使う。

F.interpolate

F.interpolate(input, size=None, scale_factor=None, mode='nearest', align_corners=None)

modeは以下の6つ

'nearest', 
'linear', 
'bilinear', 
'bicubic', 
'trilinear', 
'area'

画像バッチである4階テンソルに使えるのは以下の4つ

'nearest', 'bilinear', 'bicubic', 'area'

使い方

x = F.interpolate(x, (4,4), mode='nearest')

x = F.interpolate(x, (4,4), mode='bicubic', align_corners=False)

x = F.interpolate(x, (4,4), mode='bilinear', align_corners=False)

x = F.interpolate(x, (4,4), mode='area')

出力の例

import torch
from torch.nn import functional as F
t = torch.arange(32).reshape(2,1,4,4)
t = t.type(torch.float32)
t2 = F.interpolate(t, (2,2), mode='area') # (2,1,2,2)

以下がtの中身

tensor([[[[ 0.,  1.,  2.,  3.],
          [ 4.,  5.,  6.,  7.],
          [ 8.,  9., 10., 11.],
          [12., 13., 14., 15.]]],

        [[[16., 17., 18., 19.],
          [20., 21., 22., 23.],
          [24., 25., 26., 27.],
          [28., 29., 30., 31.]]]])

以下がt2の中身

tensor([[[[ 2.5000,  4.5000],
          [10.5000, 12.5000]]],

        [[[18.5000, 20.5000],
          [26.5000, 28.5000]]]])

nn.Upsample

これもF.interpolateと同じ

nn.Upsample(size=None, scale_factor=None, mode='nearest', align_corners=None,)

使い方

import torch
from torch import nn

x = torch.randn(2,2)
m = nn.Upsample((4,4), mode='bicubic')
x = m(x)

F.adaptive_avg_pool2d

F.adaptive_avg_pool2d(input, output_size)

pytorchのモデル中でのリサイズはadaptive average pooling 2d を使ってやることが多い印象。小さな誤差はconv2dの学習で吸収できるので、areaなのかcubicなのかを区別する必要はないんだろう。

new_h, new_w = 5,5
new_t = F.adaptive_avg_pool2d(t, (new_h, new_w))

と書けば、テンソルの縦横をリサイズできて勾配も逆伝搬できる。

import torch
from torch.nn import functional as F
t = torch.arange(32).reshape(2,1,4,4)
t = t.type(torch.float32)
t2 = F.adaptive_avg_pool2d(t, (2,2)) # (2,1,2,2)

結果はinterpolateでmode='area'を選んだ場合と同じになる。

もちろん F.adaptive_max_pool2d という関数もある。

new_h, new_w = 5,5
new_t = F.adaptive_max_pool2d(t, (new_h, new_w))

modelの中で使う場合のサンプルコード

pytorch の model 内部で定義して使う場合のサンプルコードも載せておく。

import torch
from torch import nn
from torch.nn import functional as F
from torch import optim


# networkの定義
class Net(nn.Module):
    def __init__(self):
        super().__init__() # python3系ではクラス名は省略してよい
        self.conv2d = nn.Conv2d(1, 2, (3,3), stride=1, padding=1)

    def forward_1(self, x):
        x = self.conv2d(x)
        x = F.interpolate(x, (2,2), mode='area')
        return x

    def forward_2(self, x):
        x = self.conv2d(x)
        x = F.adaptive_avg_pool2d(x, (2,2))
        return x


# ここからmain
if __name__ == '__main__':

    # GPUかCPUかを自動設定
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    model = Net().to(device)
    opt = optim.Adam(model.parameters())
    x = torch.randn(2,1,4,4).to(device)
    y = torch.ones(2,2,2,2).type(torch.float32).to(device)

    for i in range(1000):
        y1 = model.forward_1(x)
        y2 = model.forward_2(x)
        loss1 = torch.sum((y-y1)**2)
        loss2 = torch.sum((y-y2)**2)
        loss = loss1 + loss2

        opt.zero_grad()
        loss.backward()
        opt.step()

        if i%100==0: # 100回に1回lossを表示
            print('loss=%.5f' % loss.item())