pytorchがトレーニング済みモデルを保存すると、トレーニングプロセスで使用されたデバイス番号(GPUカード番号cuda:0
などcpu
)も保存されます。pytorchが履歴モデルをリロードすると、モデルはデフォルトで、トレーニング中のデバイスカード番号に従って対応するカード番号にウェイトをロードします。
ただし、テストプロセスとトレーニングプロセスの機器の状態に一貫性がない場合があります。
たとえば、ホストAには4つのGPUカードがありcuda:3
、モデルのトレーニングと保存に使用します。
テストでは、顧客のホストBでモデルを実行する必要がありますが、Bは1つのgpuカードのみをホストしますcuda:0
。
モデルがデフォルトの方法でロードされている場合、pytorchはGPUデバイスが見つからないこと、またはその他のエラーを報告します。
この時点で、ロード時に変換を行って、torch.loadのgpuデバイスのマッピング方法を指定する必要があります。
#模型的保存和加载
import torch
import json
import logger_wrappers
import os
destination_device='cuda:0' #目标设备
model_CKPT = torch.load(path,map_location={
'cuda:0':destination_device,'cuda:1':destination_device,'cuda:2':destination_device,'cuda:3':destination_device})
map_location
これはdictであり、その中のキーはストレージモデルのソースデバイス番号を表し、値は宛先デバイス番号を表します。
{ :}履歴モデルに最初に配置されたすべての重みをデバイスにロードすることを意味します。{ :}履歴モデルに最初に配置されたすべての重みをデバイスにロードすることを意味します。
cuda:3
cuda:0
cuda:3
cuda:0
cuda:3
cpu
cuda:3
cpu