PyTorch 是一个开源的机器学习库,由 Facebook 的 AI 研究团队开发。它提供了灵活的深度学习框架,适用于各种深度学习模型的研究和开发。
特点
- 动态计算图:PyTorch 使用动态计算图,这使得模型构建和调试更加直观和灵活。
- 易于使用:PyTorch 提供了简单易用的 API,使得开发者可以快速上手。
- 强大的社区支持:PyTorch 拥有一个活跃的社区,提供了大量的教程和资源。
安装
您可以通过以下命令安装 PyTorch:
pip install torch torchvision torchaudio
快速入门
- 创建一个神经网络:
import torch
import torch.nn as nn
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()
print(net)
- 训练网络:
import torch.optim as optim
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
for epoch in range(2): # 训练两个周期
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = net(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if i % 2000 == 1999: # 每2000个数据打印一次
print(f'[{epoch + 1}, {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
print('Finished Training')
更多资源
您可以在本站的 PyTorch 教程 中找到更多关于 PyTorch 的资源。
[
[