Cuando pytorch guarda el modelo entrenado, será también ahorrar el número de dispositivo utilizado en el proceso de formación (como número de tarjeta GPU cuda:0, cpu). Cuando pytorch recarga el modelo histórico, el modelo predeterminado carga el peso en el número de tarjeta correspondiente de acuerdo con el número de tarjeta del dispositivo durante el entrenamiento.
Sin embargo, a veces el proceso de prueba y la condición del equipo del proceso de capacitación son inconsistentes.
Por ejemplo, el host A tiene cuatro tarjetas GPU, y luego usamos para cuda:3entrenar el modelo y guardar el modelo.
En la prueba, hay que ejecutar el modelo en el del cliente host B, pero B anfitriones tarjeta de una sola GPU: cuda:0.
Si el modelo se carga de la forma predeterminada, pytorch informará que no puede encontrar el dispositivo gpu o algunos otros errores.
En este punto, debe realizar una transformación al cargar para especificar el método de mapeo del dispositivo gpu para torch.load:
#模型的保存和加载
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_locationEs un dict, la clave representa el número de dispositivo de origen del modelo de almacenamiento y el valor representa el número de dispositivo de destino.
{ : } Significa cargar todos los pesos colocados originalmente en el modelo histórico en el dispositivo. { : } Significa cargar todos los pesos colocados originalmente en el modelo histórico en el dispositivo.
cuda:3cuda:0cuda:3cuda:0
cuda:3cpucuda:3cpu