PyTorch Tensorを確実にNumpy Arrayに変換する

タグ：python machine_learning

結論

Tensor.detach().cpu().numpy() とする。

  In [1]: import torch
  In [2]: import numpy as np
  In [3]: a = torch.Tensor([[1, 2], [3, 4]])
  In [4]: a.detach().cpu().numpy()
  Out[5]:
  array([[1., 2.],
         [3., 4.]], dtype=float32)

とすれば（ほぼ）確実に変換が成功する。解説：PyTorch TensorとNumpy Arrayはどちらも多次元配列のためのデータ型なので自由に変換できていいのでは？と思うが実はTensorには機械学習のための色々な追加機能があり、
これを無効化してやらないとNumPy側で対応していない部分が
残ってしまい変換できないのが問題だ。

失敗例1：テンソルがGPUにある

to("cuda")でテンソルをGPUに移しておいて高速で計算できるようにできるが、 NumPyはGPUの機能がないため変換に失敗（ちなみにNumPyのGPU版のCuPyというのもあるらしい）

    In [7]: a = torch.Tensor([[1, 2], [3, 4]]).to("cuda")
    In [8]: a.numpy()
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
 in 
----> 1 a.numpy()

TypeError: can't convert cuda:0 device type tensor to numpy.
Use Tensor.cpu() to copy the tensor to host memory first.

ちなみにTensor.cpu()の代わりにTensor.to("cpu")と書いても同じである。

失敗例1：テンソルがgradを持っている

Gradとはニューラルネットなどの学習に使うテンソルの勾配の情報である。
これは機械学習でよく使うのでPyTorchではテンソルとセットで持っているようになっているが、当然NumPyでサポートされていないので変換を妨げる。
detach()でgradを持たないテンソルのコピーを作ることができ、これはNumPyに変換できる。

    In [10]: a = torch.nn.Parameter(torch.Tensor([[1, 2], [3, 4]]))  # Parameterはgradを持っている
    In [11]: a.numpy()
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
 in 
----> 1 a.numpy()

RuntimeError: Can't call numpy() on Tensor that requires grad.
Use tensor.detach().numpy() instead.

対策

Gradをなくしてくれるdetach()とcpuにデータを移すcpu()を必ず呼ぶようにする。
これらはもともとgradがなかったりデータがcpuにあっても何もしないため、呼んでおいて損はない。

    In [12]: a = torch.nn.Parameter(torch.Tensor([[1, 2], [3, 4]])).to("cuda")

  In [13]: a.detach().cpu().numpy()
  Out[13]:  #成功
  array([[1., 2.],
         [3., 4.]], dtype=float32)

Clone/copyは必要ないならしなくてもよい

ネット上ではTensor.detach().cpu().clone().numpy()だとか Tensor.detach().cpu().numpy().copy()のようにテンソルの複製までまとめてやるのを推奨している記事もあることに気づいた。
確かにこうしておけばメモリも確保しなおしてくれるので安全であるが、特にテンソルが大きいときはメモリの無駄にもなるので個人的にはclone/copyは必要な時だけするのが良いと思う。

Tensor.detach().cpu().numpy()だけであると元のテンソルと新しいarray の間でデータを保持するメモリ領域が共有されることがある。（なおこれが起こるのはもとのテンソルもcpu側にあるときだけである。to("cpu")はgpuメモリ上のデータに対してはコピーを走らせるので、このような共通参照による問題は起こらない）
これ問題となるのは変換の後でもとのテンソルに書き換えが起こるケースだ

    In [11]: a = torch.Tensor([[1, 2], [3, 4]])

    In [12]: b = a.detach().cpu().numpy()

    In [13]: a
    Out[13]:
    tensor([[1., 2.],
            [3., 4.]])

    In [14]: b
    Out[14]:
    array([[1., 2.],
          [3., 4.]], dtype=float32)

    In [15]: a[0, 0] = 5

    In [16]: b
    Out[16]:  # こちらも変更されていることに注意、これをよく見落とす
    array([[5., 2.],
          [3., 4.]], dtype=float32)

ここでcopyが入っているとaの変更はbに波及しない。メモリ共有が起こらないためだ。

    In [11]: a = torch.Tensor([[1, 2], [3, 4]])

    In [12]: b = a.detach().cpu().numpy().copy()

    In [13]: a
    Out[13]:
    tensor([[1., 2.],
            [3., 4.]])

    In [14]: b
    Out[14]:
    array([[1., 2.],
          [3., 4.]], dtype=float32)

    In [15]: a[0, 0] = 5

    In [23]: b
    Out[23]:  # こちらは変更なし
    array([[1., 2.],
           [3., 4.]], dtype=float32)
    
    In [24]: a
    Out[24]:  # こちらだけ確かに書き変わっている
    tensor([[5., 2.],
            [3., 4.]])

動作確認したバージョン

Python 3.10.5
torch==1.13.1+cu116
numpy==3.10.5

PyTorch Tensorを確実にNumpy Arrayに変換する

結論

失敗例1：テンソルがGPUにある

失敗例1：テンソルがgradを持っている

対策

Clone/copyは必要ないならしなくてもよい

動作確認したバージョン

おすすめ記事