import torch

# 创建一个张量
tensor = torch.tensor([1, 2, 3, 4, 5])

# 访问单个元素
print(tensor[0])  # 输出：1

# 切片操作
print(tensor[1:4])  # 输出：tensor([2, 3, 4])

# 修改元素的值
tensor[0] = 10
print(tensor)  # 输出：tensor([10,  2,  3,  4,  5])

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms

# 设置随机种子，以确保结果的可重复性
torch.manual_seed(0)

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize((0.1307,), (0.3081,))  # 标准化图像
])

# 加载MNIST数据集
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)  # 训练数据集
test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)  # 测试数据集

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)  # 训练数据加载器
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)  # 测试数据加载器

# 定义神经网络模型（LeNet）
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 20, 5, 1)  # 第一个卷积层
        self.conv2 = torch.nn.Conv2d(20, 50, 5, 1)  # 第二个卷积层
        self.fc1 = torch.nn.Linear(4*4*50, 500)  # 第一个全连接层
        self.fc2 = torch.nn.Linear(500, 10)  # 第二个全连接层（输出层）

    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))  # 第一个卷积层使用ReLU激活函数
        x = torch.nn.functional.max_pool2d(x, 2, 2)  # 最大池化层
        x = torch.nn.functional.relu(self.conv2(x))  # 第二个卷积层使用ReLU激活函数
        x = torch.nn.functional.max_pool2d(x, 2, 2)  # 最大池化层
        x = x.view(-1, 4*4*50)  # 将特征图展平
        x = torch.nn.functional.relu(self.fc1(x))  # 第一个全连接层使用ReLU激活函数
        x = self.fc2(x)  # 输出层
        return torch.nn.functional.log_softmax(x, dim=1)  # 使用log_softmax作为输出

# 创建神经网络模型实例
model = Net()

# 定义优化器和损失函数
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)  # 使用随机梯度下降优化器
criterion = torch.nn.CrossEntropyLoss()  # 交叉熵损失函数

# 训练模型
def train(epoch):
    model.train()  # 设置模型为训练模式
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()  # 清除梯度
        output = model(data)  # 进行前向传播
        loss = criterion(output, target)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        if batch_idx % 100 == 0:
            print('训练 Epoch: {} [{}/{} ({:.0f}%)]\t损失: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

# 测试模型
def test():
    model.eval()  # 设置模型为评估模式
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\n测试集: 平均损失: {:.4f}, 准确率: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

# 进行模型训练和测试
for epoch in range(1, 11):
    train(epoch)  # 训练模型
    test()  # 测试模型