pytorch .to(device)用法和原理

PyTorch中，.to(device) 是一个非常重要的方法，用于将张量、模型等对象移动到指定的设备（如CPU或GPU）。

import torch
import torch.nn as nn

# 检查可用设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"使用设备: {device}")

# 对于张量
tensor_cpu = torch.tensor([1, 2, 3])
tensor_gpu = tensor_cpu.to(device)  # 移动到GPU（如果可用）
print(f"原始张量设备: {tensor_cpu.device}")
print(f"移动后张量设备: {tensor_gpu.device}")

# 对于模型
model = nn.Linear(10, 5)
model = model.to(device)  # 将模型移动到指定设备
print(f"模型设备: {next(model.parameters()).device}")

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear1 = nn.Linear(10, 50)
        self.relu = nn.ReLU()
        self.linear2 = nn.Linear(50, 1)
    
    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

def train_model():
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"训练设备: {device}")
    
    # 初始化模型、数据、优化器
    model = SimpleModel()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.MSELoss()
    
    # 将模型移动到设备
    model = model.to(device)
    
    # 生成示例数据
    x_train = torch.randn(100, 10)  # 在CPU上创建数据
    y_train = torch.randn(100, 1)
    
    # 将数据移动到相同设备
    x_train = x_train.to(device)
    y_train = y_train.to(device)
    
    # 训练循环
    model.train()
    for epoch in range(5):
        optimizer.zero_grad()
        
        # 前向传播（自动在正确设备上计算）
        outputs = model(x_train)
        loss = criterion(outputs, y_train)
        
        # 反向传播
        loss.backward()
        optimizer.step()
        
        print(f'Epoch [{epoch+1}/5], Loss: {loss.item():.4f}')

if __name__ == "__main__":
    train_model()

 

模型移动到GPU

model = MyModel()
model = model.to(device)  # 将模型参数和缓冲区移到GPU

作用：

• 将模型的所有参数（weights, biases）和缓冲区移到指定设备

• 确保前向传播和反向传播的计算在GPU上执行

• 只需要执行一次，在训练开始前完成

 

数据移动到GPU

# 单个数据
x = torch.tensor([1, 2, 3])
x = x.to(device)

# 批量数据（通常在训练循环中）
for batch in dataloader:
    inputs, labels = batch
    inputs = inputs.to(device)
    labels = labels.to(device)

作用：

• 将输入数据和标签移到与模型相同的设备

• 需要在每个batch处理时执行，因为数据是不断加载的

 

注意：

1.设备一致性

模型和数据必须在同一设备上，否则会报错

# 错误：模型在GPU，数据在CPU
RuntimeError: Expected all tensors to be on the same device

核心原则：模型迁移是一次性的，数据迁移是持续性的，两者都必须执行且设备要一致。

2.torch.optim.Adam 

优化器不需要显式指定到GPU上，会自动管理

optimizer = torch.optim.Adam(model.parameters()) 会获取模型参数的引用

优化器的内部状态（如动量缓冲区）会自动创建在与参数相同的设备上

optimizer.step() 在原地更新参数，保持设备一致性

# ✅ 正确的做法
model = nn.Linear(10, 2)
model = model.to(device)  # 先移动模型
optimizer = torch.optim.Adam(model.parameters())  # 优化器自动使用GPU