検索
  • Kunita

【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)))

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







閲覧数:2,218回0件のコメント

最新記事

すべて表示

Unityのプロジェクトを外付けHDDへ移動させたところ、下記のエラーが出るようになった。 Microsoft (R) Visual C# Compiler version 2.9.1.65535 (9d34608e) Copyright (C) Microsoft Corporation. All rights reserved. error CS2012: Cannot open ***/U

突然ですが、皆さん!AIに興味ありますか〜? ゲーム開発に興味がある人は、「対戦型ゲームで敵AIの作成」なんて課題を抱えている方も多いかも知れませんね。 ゲームを面白くするためには、AIの存在は必須かも知れません。 実は、Unityには、AIの作成を補助してくれる機能があるんです。しかも、近年、話題の囲碁AIみたいに製作者が特に何もしなくても、色んなことを勝手に覚えて強いNPCを制作することができ