pytorchが履歴モデルをロードするときにGPUカード番号を変更します。map_location設定

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:3cuda:0cuda:3cuda:0
cuda:3cpucuda:3cpu

おすすめ

転載: blog.csdn.net/jmh1996/article/details/111041108