PyTorch框架的使用主要包括以下步驟:
以下是一個簡單的PyTorch使用示例,用于實現一個線性回歸模型:
import torch
import torch.nn as nn
import torch.optim as optim
# 定義數據集
x_train = torch.Tensor([[1.0], [2.0], [3.0], [4.0]])
y_train = torch.Tensor([[2.1], [3.9], [5.3], [6.7]])
# 定義模型
class LinearRegressionModel(nn.Module):
def __init__(self, input_dim, output_dim):
super(LinearRegressionModel, self).__init__()
self.linear = nn.Linear(input_dim, output_dim)
def forward(self, x):
out = self.linear(x)
return out
# 實例化模型、損失函數和優化器
input_dim = 1
output_dim = 1
model = LinearRegressionModel(input_dim, output_dim)
criterion = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)
# 訓練模型
num_epochs = 1000
for epoch in range(num_epochs):
# 前向傳播
outputs = model(x_train)
loss = criterion(outputs, y_train)
# 反向傳播和優化
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (epoch+1) % 100 == 0:
print ('Epoch [{}/{}], Loss: {:.4f}'.format(epoch+1, num_epochs, loss.item()))
# 測試模型
with torch.no_grad():
test_input = torch.Tensor([[5.0]])
test_output = model(test_input)
print('Test Output: {}'.format(test_output))
在這個示例中,我們首先定義了一個簡單的線性回歸模型,并使用隨機生成的數據集進行訓練和測試。然后,我們使用PyTorch提供的訓練循環來訓練模型,并在每個epoch結束后輸出損失值。最后,我們在測試數據集上評估模型的性能,并輸出測試輸出。