明月深度学习实践002：关于模型训练你应该关注的内容

# 判断cuda是否可用
torch.cuda.is_available()

# 定义device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 在数据准备好之后，应该统一转移到device
images = images.to(device)

# 查看显存的占用，除以2个1024，单位位MB
torch.cuda.max_memory_allocated()//(1024*1024)

# 释放显存占用
# empty_cache()来释放所有未使用的缓存的内存，以便其它 GPU 应用能够使用. 但是并不能释放 tensors 所占用的 GPU 显存
torch.cuda.empty_cache()

        self.conv1 = nn.Sequential(     # input: 28*28*1
            # filter 大小 5×5，filter 深度（个数）为 6，padding为2，卷积步长s=1
            # 输出矩阵大小为 28×28×6
            nn.Conv2d(1, 6, 5, 1, 2),
            nn.ReLU(),
            # filter 大小 2×2（即 f=2），步长 s=2，no padding
            # 输出：14*14*6
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.conv2 = nn.Sequential(
            # filter 大小 5×5，filter 个数为 16，padding 为 0， 卷积步长 s=1
            # 输出矩阵大小为 10×10×16
            nn.Conv2d(6, 16, 5),
            nn.ReLU(),
            # 输出：5*5*16
            nn.MaxPool2d(2, 2)
        )
        # 全连接层
        self.fc1 = nn.Sequential(
            nn.Linear(16 * 5 * 5, 120),
            nn.ReLU()
        )
        # 全连接层
        self.fc2 = nn.Sequential(
            nn.Linear(120, 84),
            nn.ReLU()
        )
        # 全连接层，输出神经元数量为 10，代表 0～9 十个数字类别
        self.fc3 = nn.Linear(84, 10)

import time
import torch.optim as optim

EPOCH = 10        # 遍历数据集次数
LR = 0.001        # 学习率

# 设置mlflow及记录参数
mlflow.set_experiment('LeNet测试实验')
mlflow.start_run()
mlflow.log_param('EPOCH', EPOCH)
mlflow.log_param('LR', LR)

# 定义是否使用GPU
print("cuda" if torch.cuda.is_available() else "cpu")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 定义损失函数loss function 和优化方式（采用SGD）
net = LeNet().to(device)
criterion = nn.CrossEntropyLoss()  # 交叉熵损失函数，通常用于多分类问题上
optimizer = optim.SGD(net.parameters(), lr=LR, momentum=0.9)

start = time.time()
for epoch in range(EPOCH):
    sum_loss = 0.0
    # 数据读取
    for i, data in enumerate(trainloader):
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)

        # 梯度清零
        optimizer.zero_grad()

        # forward + backward
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        # 每训练100个batch打印一次平均loss
        sum_loss += loss.item()
        if i % 100 == 99:
            print('[%d, %d] loss: %.03f' % (epoch + 1, i + 1, sum_loss / 100))
            sum_loss = 0.0

    # 每跑完一次epoch测试一下准确率
    with torch.no_grad():
        correct = 0
        total = 0
        for data in testloader:
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = net(images)
            # 取得分最高的那个类
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum()
        print('第%d个epoch的识别准确率为：%.2f%%' % (epoch + 1, (100 * int(correct) / total)))
        mlflow.log_metric('Accuracy', 100 * int(correct) / total, step=epoch)

    # 保存模型
    torch.save(net.state_dict(), 'models/lenet/lenet_%03d.pth' % epoch)
    
# 记录运行时间与显存占用
# 这两个值对于非常重要，记录下来
mlflow.log_metric('gpu', torch.cuda.max_memory_allocated()//(1024*1024))
mlflow.log_metric('run_time', time.time()-start)
# 结束记录
mlflow.end_run()

# 释放显存占用
torch.cuda.empty_cache()

print(torch.max(img_conv1), torch.min(img_conv1))
max_val = torch.max(img_conv1).to('cpu')
for i in range(img_conv1.shape[1]):
    _img = img_conv1[0][i]
    _img = _img.to('cpu').clone()
    _img = (_img/max_val)*255
    img_pil = transforms.ToPILImage()(_img)
    img_pil = img_pil.resize((112, 112))
    display(img_pil)