【pytorch】GPUで訓練させたモデルをCPUで読み込もうとしたお話

タイトル通り。

PytorchでStyleGANを試していたのだが、完成したモデルををCPUで読み込んで動かそうとしたらうまくいかなかったので、頑張って直した話。

そもそもGPUで学習させたんだから、同じくGPUでやれよ、という話なのだが、「GoogleColaboratoryで作ったモデルをローカルで動かしたかった」というのが最大の動機。

それに、CPUだけで処理できると、AI学習済みモデルをデプロイしたwebアプリケーション(例えば、GANで画像生成できるサイトとかを用意した時とか)で、貸サーバーも

無料 or 格安のものが使えるなぁ。。なんて思ったので、頑張ってやってみました。

実行環境

MacOS 11.2.1

Python 3.8

torch 1.7.0

Q. そもそもどんな問題が起こるのか？

A. 下記エラー文が出力される。

RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU.

読み込み時のエラーなので、まず、モデルの読み込みを行うスクリプトが記載されているファイルに下記クラスを定義・追記。

import torch
import io
        
class CPU_Unpickler(pickle.Unpickler):
    def find_class(self, module, name):
        if module == 'torch.storage' and name == '_load_from_bytes':
            return lambda b: torch.load(io.BytesIO(b), map_location='cpu')
        else: return super().find_class(module, name)

これにより、gpu定義モデルをgpuで読み込めるクラスを定義完了。続けて、モデル読み込み部分を以下のように修正。

if torch.cuda.is_available():
    with open(restore_model_file, 'rb') as file:
        ret = pickle.load(file, encoding='latin1')
else:
    with open(restore_model_file, 'rb') as file:
        ret = CPU_Unpickler(file).load()

cudaが利用できない場合の処理をif文で記述しておく。

多分、ここまでやればモデルの読み込みができる筈。

ただ、自分の使っていたソースコードでは下記TypeErrorが出たため、さらに修正した。

TypeError: type torch.cuda.FloatTensor not available. Torch not compiled with CUDA enabled.

これは、スクリプトのどこかでtorch.cuda.FloatTensorが使われているということ。多分、クラス継承時かどこかやらかしている。torch.Tensorを使うように修正している必要がある。エラーログから該当部分を見つけ、下記のようにtorch.Tensorを使うように修正

if torch.cuda.is_available():
    noise = self.Tensor(np.random.normal(loc=0, scale=1, size=(batch_size, 1, height, width)))
else:
    noise = torch.Tensor(np.random.normal(loc=0, scale=1, size=(batch_size, 1, height, width)))

以上により、問題が修正され、スクリプトが正常に動作するようになった。