【Colab】基本操作【LeNet】【MNIST】训练测试

Colab 官网初始界面

1.介绍

英伟达官网
谷歌将原来K80换成了T4

2.查看基本配置

2.1查看pytorch版本

import torch

print(torch)

2.2查看是否可以使用cuda

print(torch.cuda.device_count())

修改->笔记本设置->GPU

2.3查看显卡配置

!nvidia-smi
//注意英文感叹号

3.挂载

31.挂载谷歌云盘

Colab的运行原理实际上就是给你分配一台远程的带GPU的主机，所以它的原始路径不是你的谷歌云盘（也就是你的代码文件）所在的路径。所以第一步我们先要把谷歌云盘挂载带到那台远程主机上：

from google.colab import drive
drive.mount("/content/drive")

登录谷歌账号并将验证码粘到框中

3.2更改运行目录

import sys
sys.path.append("/content/drive/MyDrive/Colab Notebooks/LeNet_MNIST_train_test")

下面是我的目录结构

4.训练

【LeNet】+【MNIST】

import torchvision
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import transforms
import time

class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(1, 6, kernel_size=5),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(6, 16, kernel_size=5),
            nn.MaxPool2d(kernel_size=2),
            nn.Flatten(),
            nn.Linear(16*4*4, 120),
            nn.Linear(120, 84),
            nn.Linear(84, 10)
        )

    def forward(self, x):
        output = self.model(x)
        return output

train_datasets = torchvision.datasets.MNIST(
    root = r'../data',
    download=True,
    train=True,
    transform=transforms.ToTensor()
)
train_dataloader = DataLoader(
    dataset=train_datasets,
    batch_size=64
)

test_datasets = torchvision.datasets.MNIST(
    root = r'../data',
    train=False,
    download=True,
    transform=transforms.ToTensor()
)
test_dataloader = DataLoader(
    dataset=test_datasets,
    batch_size=64
)

train_datasets_size = len(train_datasets)
test_datasets_size = len(test_datasets)
print("训练集数量为：{}".format(train_datasets_size))
print("测试集数量为：{}".format(test_datasets_size))

runing_mode = "gpu" # cpu,gpu, gpus
if runing_mode == "gpu" and torch.cuda.is_available():
    print("use cuda")
    device = torch.device("cuda")
else:
    print("use cpu")
    device = torch.device("cpu")

model = LeNet()
model.to(device)

loss_fn = nn.CrossEntropyLoss()
loss_fn.to(device)
learning_rate = 1e-2
optim = torch.optim.SGD(model.parameters(), lr=learning_rate)

epoch = 10
train_step, test_step = 0, 0
for i in range(epoch):
    print("~~~~~~~~~~~~第{}轮训练开始~~~~~~~~~~~".format(i+1))
    start = time.time()
    model.train()
    for data in train_dataloader:
        imgs, targets = data
        imgs, targets = imgs.to(device), targets.to(device)
        output = model(imgs)
        loss = loss_fn(output, targets)

        optim.zero_grad()
        loss.backward()
        optim.step()

        train_step += 1
        if train_step % 200 == 0:
            print("第{}次训练，loss={:.3f}".format(train_step, loss.item()))
    #
    model.eval()
    with torch.no_grad():
        test_loss, true_num = 0, 0
        for data in test_dataloader:
            imgs, targets = data
            imgs, targets = imgs.to(device), targets.to(device)
            output = model(imgs)
            test_loss += loss_fn(output, targets)
            true_num += (output.argmax(1) == targets).sum()
    end = time.time()
    print("第{}轮测试集上的loss:{:.3f}, 正确率为:{:.3f}%,耗时:{:.3f}".format(test_step+1, test_loss.item(), 100 * true_num / test_datasets_size, end-start))
    test_step += 1

运行结果：
1.CPU

2.GPU

5.Reference

《Colab使用训练指南》 坚强的羊脂球