010 CNN基础

代码详解

batch_size = x.size(0) 这一行代码是用来获取当前批次（batch）的数据量。在 PyTorch 中，张量 x 的维度通常按照 [batch_size, channel, height, width] 的顺序排列，其中 batch_size 表示一批次中包含的样本数量。
x = x.view(batch_size, -1) 这一行中用于将卷积层的输出展平成一个一维向量，其中 -1 表示自动计算除 batch_size 外的其他维度，即这里的 320。这样就可以将展平后的向量输入到全连接层中了。

import torch
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import torch.nn.functional as F
import torch.optim as optim

# prepare dataset

# 把原始图像转换为图像张量，C*W*H  然后对数据进行归一化，类比正态分布转化为标准形式
batch_size = 64
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,),(0.3081,))])

train_dataset = datasets.MNIST(root='../dataset/mnist/', train=True, download=True,transform=transform)
train_loader = DataLoader(train_dataset, shuffle=True,batch_size=batch_size)
test_dataset = datasets.MNIST(root='../dataset/mnist/', train=False, download=True,transform=transform)
test_loader = DataLoader(test_dataset, shuffle=False, batch_size=batch_size)

# design model using class

class Net(torch.nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 10, kernel_size=5)       # 卷积核的大小是5*5
        self.conv2 = torch.nn.Conv2d(10, 20, kernel_size=5)
        self.pooling = torch.nn.MaxPool2d(2)    # 提供2维的最大池化操作
        # 定义全连接层
        self.fc = torch.nn.Linear(320,10)  # 输入大小为320（假设是展平后的向量长度），输出大小为10（分类任务中的类别数量）


    def forward(self,x):
        # flatten
        batch_size = x.size(0)   #获取当前批次的数据量
        x = F.relu(self.pooling(self.conv1(x)))
        x = F.relu(self.pooling(self.conv2(x)))
        x = x.view(batch_size, -1)
        x = self.fc(x)
        return x


model = Net()
# 调用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)


criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)


def train(epoch):
    running_loss = 0.0
    for batch_idx, data in enumerate(train_loader, 0):
        inputs, target = data
        inputs, target = inputs.to(device), target.to(device)        #GPU
        outputs = model(inputs)
        loss = criterion(outputs, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        if batch_idx % 300 == 299:
            print('[%d,%5d]loss:%.3f' % (epoch + 1, batch_idx + 1, running_loss / 300))


def Test():
    correct = 0
    total = 0
    with torch.no_grad():             # 禁用梯度计算。在测试集上不需要反向传播，禁用梯度计算减少内存消耗
        for data in test_loader:
            images, labels =data
            images, labels = images.to(device), labels.to(device)          # GPU
            outputs = model(images)
            _, predicted = torch.max(outputs.data, dim=1)           # _说明不关注类别本身。只关注索引标签，在这里把标签预测的标签记为：predicted
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    print('accuracy on test set: %d %%' % (100*correct/total))


if __name__ == '__main__':
    for epoch in range(10):
        train(epoch)
        Test()