如何在PyTorch中使用GPU加速計算

在PyTorch中使用GPU加速計算可以通過以下步驟實現：

檢查是否有可用的GPU設備：

import torch

if torch.cuda.is_available():
    print("GPU is available!")
else:
    print("GPU is not available.")

將Tensor對象移動到GPU設備上：

# 創建一個Tensor對象
x = torch.randn(3, 3)

# 將Tensor對象移動到GPU設備上
x = x.to('cuda')

在定義模型時指定使用GPU設備：

import torch.nn as nn

# 定義一個簡單的神經網絡模型
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 5)
        self.fc2 = nn.Linear(5, 2)
        
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 創建模型實例并將其移動到GPU設備上
model = SimpleNN()
model = model.to('cuda')

在訓練過程中使用GPU設備：

# 將輸入數據和標簽移動到GPU設備上
inputs = inputs.to('cuda')
labels = labels.to('cuda')

# 使用GPU設備進行前向傳播和反向傳播
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()

通過以上步驟，可以在PyTorch中使用GPU加速計算，提高訓練模型的速度和效率。

亚洲激情专区-91九色丨porny丨老师-久久久久久久女国产乱让韩-国产精品午夜小视频观看

最新問答

相關標簽