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