torch.load()导入模型pt文件

在PyTorch中，`torch.load()`函数用于从.pt或.pth文件中加载预训练模型。当你想从GPU上保存的模型文件导入到CPU或另一个GPU时，需要特别处理以适应不同的硬件环境。以下是基本步骤：

1. **导入模型**:

   # 假设你有一个保存的模型路径model_path
   model = torch.load(model_path)

2. **加载到CPU（默认）**:
如果不需要在GPU上运行，`torch.load()`会自动将模型从原设备（通常是GPU）迁移到CPU。

3. **加载到特定GPU**:
**如果你的设备ID为0**:

   if torch.cuda.is_available():
       device = torch.device('cuda:0')
       model.load_state_dict(torch.load(model_path, map_location=device))
       model = model.to(device)

**如果设备ID不是0**:

   if torch.cuda.is_available():
       device = torch.device(f'cuda:{your_device_id}')
       model.load_state_dict(torch.load(model_path, map_location=device))
       model = model.to(device)

4. **加载数据到GPU**:
当处理模型输入时，确保它们也在相应的GPU上:

   input_data = data.to(device)

请注意，确保在执行上述操作之前，`torch`库已经安装并且支持CUDA。另外，如果你尝试在非GPU环境中运行这段代码，`map_location`选项将不起作用，因为GPU资源不可用。