PyTorch 是一个开源的机器学习库,由 Facebook 的 AI 研究团队开发。它提供了灵活的深度学习框架,被广泛应用于计算机视觉、自然语言处理等领域。
特点
- 动态计算图:PyTorch 使用动态计算图,这使得调试和原型设计变得更加容易。
- 易于使用:PyTorch 的 API 设计简洁,易于上手。
- 强大的社区支持:PyTorch 拥有庞大的社区,提供了大量的教程和资源。
快速开始
要开始使用 PyTorch,首先需要安装它。您可以通过以下命令安装:
pip install torch torchvision
安装完成后,您可以创建一个简单的神经网络:
import torch
import torch.nn as nn
import torch.optim as optim
# 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 6, 3)
self.conv2 = nn.Conv2d(6, 16, 3)
self.fc1 = nn.Linear(16 * 6 * 6, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 10)
def forward(self, x):
x = torch.relu(self.conv1(x))
x = torch.max_pool2d(x, 2, 2)
x = torch.relu(self.conv2(x))
x = torch.max_pool2d(x, 2, 2)
x = x.view(-1, self.num_flat_features(x))
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
def num_flat_features(self, x):
size = x.size()[1:] # 除批量大小外的所有维度
num_features = 1
for s in size:
num_features *= s
return num_features
net = Net()
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# 训练模型
for epoch in range(2): # 训练两个周期
for data, target in trainloader:
optimizer.zero_grad()
output = net(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
print('Finished Training')
更多关于 PyTorch 的教程和文档,请访问PyTorch 官方网站。
相关资源
