Pytorch tensor と numpy ndarray の変換

Pytorch tensor から numpy ndarray への変換とその逆変換についてまとめる。単純にtorch.from_numpy(x)とx.detach().numpy()を覚えておけばよいので、その使い方を示しておく。

すぐ使いたい場合は以下

numpy to tensor

x = torch.from_numpy(x.astype(np.float32)).clone()

tensor to numpy

x = x.to('cpu').detach().numpy().copy()

pytorchでは変数の型としてほとんどtorch.floatを使うので、基本的にはnumpyでnp.float32にキャストしてからtorch tensorへ変換する。また、torch tensor と numpy ndarray はメモリを共有しているようなので、clone()やcopy()を書いておくと書き換えの心配がなくなる。

※最近知ったんだけど、pytorch tensorもCPUメモリであれば、matplotlibでのplotやskimage.ioでの画像ファイルとしての保存ができてしまう。OpencvやPILではまだ使えないようだが、numpyに変換する必要がどの程度あるのかは疑問。

numpy ndarray から pytorch tensorへ変換

numpy ndarray から pytorch tensor へ変換するには、 numpyのうちに型をfloat32にキャストして、from_numpy()関数を使ってtensorに変換する。

import torch
import numpy as np

x_numpy = np.random.randn(10)
x_numpy = x_numpy.astype(np.float32)
x_tensor = torch.from_numpy(x_numpy).clone()

GPUメモリに移動する場合は以下

import torch

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

pytorch tensorから numpy ndarrayへ変換

torch tensorからnumpy ndarray へ変換するには、以下のようにする。(最もシンプルな書き方)

import torch

x_tensor = torch.randn(10)
x_numpy = x_tensor.to('cpu').detach().numpy().copy()

numpyは必ずCPU上のメモリを使うため、torch tensor がGPUを使っている場合は、to('cpu')で一度CPUメモリに落としてから、detach()関数を使ってデータ部分を切り離す。その後にnumpy()関数でnumpy arrayへ変換する。 最後にcopy()してtorch tensorとメモリを共有しないようにする。

注意点として、学習の途中で画像を保存するためにto('cpu')を使ってGPUからCPUに変換している場合は、以下のようにしてtensorを元のdeviceに戻すのを忘れないようにする。これを忘れると次のイテレーションでこけることがある。

import torch

x_tensor = torch.randn(10)
device = x_tensor.device
x_numpy = x_tensor.to('cpu').detach().numpy().copy()
x_tensor = x_tensor.to(device)

transforms.ToTensorと同じ処理

画像のnumpy ndarrayをCNNに入力可能な形n,c,h,wに変換するには次のようにすればよい。これはtorchvisionのtransforms.ToTensorと同じ処理になる。

from skimage import io as skio

img = skio.imread('sample.jpg')
x = torch.from_numpy(img.astype(np.float32))
x = x.unsqueeze(0).permute(0,3,1,2)

nn.CrossEntropyLoss への入力

numpy ndarrayからpytorch tensorへの変換でひとつだけ注意しなければならないのが、nn.CrossEntropyLoss への入力。この関数で定義したcriterionの正解部分へtorch.floatを入力すると、RuntimeError: expected scalar type Long but found Floatと出てしまう。この回避方法についても書いておく。

Cross entropy loss の入力はlong intでないとエラーが出て怒られる。この場合、上記のnp.float32ではなく、np.int64を使わないといけない。

import torch
import numpy as np

criterion = torch.nn.CrossEntropyLoss()

pred = torch.rand(3, 5, requires_grad=True)
real = np.random.randint(0,5,3).astype(np.float32)
real = torch.from_numpy(real)

loss = criterion(pred, real)

このように書いて実行すると、下記のようなエラーが出る

RuntimeError: expected scalar type Long but found Float

意味は「Longタイプが期待されてるのにFloatタイプが入ってます」ということなので、CrossEntropyLossの入力だけ、Long(long intのこと, torchではtorch.long, numpyではnp.int64など)にしないといけない。どこでキャストするかは自由であるが、lossへの入力直前に書くのが一番わかりやすいかと。書き方は以下

import torch
import numpy as np

criterion = torch.nn.CrossEntropyLoss()

pred = torch.rand(3, 5, requires_grad=True)
real = np.random.randint(0,5,3).astype(np.float32)
real = torch.from_numpy(real)

loss = criterion(pred, real.type(torch.long)) # realの型キャスト