1.torch安装
conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
或者pip install
2.其他库的安装
pip install numpy pandas matplotlib scikit-learn jupyterlab seaborn
3.生成 Jupyter 配置文件
jupyter notebook --generate-config
4.搜索“notebook_dir”
5.PyTorch入门示例 - MNIST数字识别
# MNIST 手写数字识别
## 1. 环境准备
# 导入必要的库
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
%matplotlib inline
# 检查 GPU 是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")
## 2. 数据加载与预处理
# 定义数据变换
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
# 加载数据集
train_dataset = datasets.MNIST('./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('./data', train=False, transform=transform)
# 创建数据加载器
batch_size = 64
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size)
# 可视化一些样本
def show_samples(data_loader):
examples = enumerate(data_loader)
batch_idx, (example_data, example_targets) = next(examples)
plt.figure(figsize=(12, 8))
for i in range(6):
plt.subplot(2, 3, i+1)
plt.imshow(example_data[0], cmap='gray')
plt.title(f"标签: {example_targets}")
plt.axis('off')
plt.tight_layout()
show_samples(train_loader)
## 3. 定义模型
# 定义神经网络模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout2d(0.25)
self.dropout2 = nn.Dropout2d(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
output = F.log_softmax(x, dim=1)
return output
# 创建模型实例并移至设备
model = Net().to(device)
print(model)
# 定义优化器和损失函数
optimizer = optim.Adam(model.parameters(), lr=0.001)
criterion = F.nll_loss
## 4. 训练模型
# 定义训练函数
def train(model, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = criterion(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print(f'训练轮次: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
f'({100. * batch_idx / len(train_loader):.0f}%)]\t损失: {loss.item():.6f}')
return loss.item()
# 定义测试函数
def test(model, test_loader):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
output = model(data)
test_loss += criterion(output, target, reduction='sum').item()
pred = output.argmax(dim=1, keepdim=True)
correct += pred.eq(target.view_as(pred)).sum().item()
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
print(f'测试集: 平均损失: {test_loss:.4f}, 准确率: {correct}/{len(test_loader.dataset)} ({accuracy:.2f}%)')
return test_loss, accuracy
# 训练模型
epochs = 5
train_losses = []
test_losses = []
test_accuracies = []
for epoch in range(1, epochs + 1):
train_loss = train(model, train_loader, optimizer, epoch)
test_loss, accuracy = test(model, test_loader)
train_losses.append(train_loss)
test_losses.append(test_loss)
test_accuracies.append(accuracy)
## 5. 可视化结果
# 绘制训练和测试损失
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(range(1, epochs + 1), train_losses, label='训练损失')
plt.plot(range(1, epochs + 1), test_losses, label='测试损失')
plt.xlabel('轮次')
plt.ylabel('损失')
plt.legend()
plt.title('训练和测试损失')
plt.subplot(1, 2, 2)
plt.plot(range(1, epochs + 1), test_accuracies, label='测试准确率')
plt.xlabel('轮次')
plt.ylabel('准确率 (%)')
plt.legend()
plt.title('测试准确率')
plt.tight_layout()
plt.show()
# 可视化模型预测
def visualize_predictions(model, test_loader):
# 获取一批数据
examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
# 获取预测结果
with torch.no_grad():
example_data = example_data.to(device)
output = model(example_data)
example_data = example_data.cpu()
# 获取预测的类别
pred = output.argmax(dim=1, keepdim=True).cpu()
# 显示结果
plt.figure(figsize=(12, 8))
for i in range(6):
plt.subplot(2, 3, i+1)
plt.imshow(example_data[0], cmap='gray')
plt.title(f'预测: {pred[0]}, 实际: {example_targets}')
plt.axis('off')
plt.tight_layout()
visualize_predictions(model, test_loader)
## 6. 保存模型
# 保存模型
torch.save(model.state_dict(), 'mnist_cnn.pt')
print("模型已保存")
# 如何加载模型
# 创建一个新的模型实例
# new_model = Net().to(device)
# 加载保存的权重
# new_model.load_state_dict(torch.load('mnist_cnn.pt'))
登录
立即注册
