Pytorch implements cifar10 classification task

0 import package

import torch
import torch.nn as nn
import torch.nn.functional as f
from torchvision.datasets import cifar
from torchvision import transforms
from torch.utils.data import DataLoader

1 Define the network structure

class Net(nn.Module):
    """网络结构"""
    def __init__(self):
        """
        初始化网络结构和损失函数
        """
        super().__init__()
        # out_size = （in_size - k + 2p）/ s +1
        self.conv1 = nn.Conv2d(3, 32, (3, 3), (1, 1), 1)
        self.conv2 = nn.Conv2d(32, 64, (3, 3), (1, 1), 1)
        self.fc1 = nn.Linear(64 * 8 * 8, 64 * 8)
        self.fc2 = nn.Linear(64 * 8, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        """正向传播"""
        x = f.relu(f.max_pool2d(f.relu(self.conv1(x)), (2, 2)))
        x = f.relu(f.max_pool2d(f.relu(self.conv2(x)), (2, 2)))
        x = x.view(-1, 64 * 8 * 8)
        x = f.relu(self.fc1(x))
        x = f.relu(self.fc2(x))
        x = f.softmax(self.fc3(x))
        return x

2 Processing data

1. Initialization parameters, network, learning rate, batch size, data storage address, loss function, optimizer
2. Data enhancement can achieve the effect of expanding the data set and improving the fitting effect
3. Process the data into a mini-batch to improve the convergence speed. Generally, the batch is set to the power of 2, such as $2^4、2^6、2^8$ etc.

class Model:
    """处理数据、训练、评估"""
    def __init__(self, net, learning_rate=0.001, batch_size=128, path='./data/'):
        """
        初始化网络结构、参数、数据存储路径
        :param net: 网络结构
        :param learning_rate: 学习率
        :param batch_size: 批量大小
        :param path: 数据存储路径
        """
        self.net = net
        self.cost = nn.CrossEntropyLoss()
        self.optimizer = torch.optim.RMSprop(self.net.parameters(), lr=learning_rate)
        self.batch_size = batch_size
        self.path = path

        # 数据增强
        self.transforms = transforms.Compose([
            transforms.CenterCrop([32, 32]),
            transforms.ToTensor(),
            transforms.Normalize((0.0, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

        # 获取数据
        self.train = cifar.CIFAR10(root=self.path, train=True, transform=self.transforms, download=True)
        self.train_loader = DataLoader(dataset=self.train, batch_size=self.batch_size, shuffle=True)
        self.test = cifar.CIFAR10(root=self.path, train=False, transform=self.transforms)
        self.test_loader = DataLoader(dataset=self.test, batch_size=self.batch_size, shuffle=True)

3 training

Model类中

For each batch of training, the number of epochs can be set.

1. set epechs
2. Divide training data and labels
3. Clear the gradient value zero_grad()
4. Training, backpropagation backward()
5. iteration parameter optimizer.step()

def train_(self, epochs=3):
    for epoch in range(epochs):
        losses = 0.0
        for i, data in enumerate(self.train_loader, 0):
            train_data, train_label = data
            self.optimizer.zero_grad()
            loss = self.cost(self.net(train_data), train_label)
            loss.backward()
            self.optimizer.step()
            losses += loss.item()
            if i % 10 == 0 and i:
                print(f'[epoch {
      
      epoch + 1}/{
      
      epochs}, {
      
      (i + 1) / len(self.train_loader) * 100:.2f}%] loss: {
      
      losses / 100:.2f}')
                losses = 0.0
    print('Train Finished!')

4 evaluation

Model类中

Evaluate the effect of the test set and write the tick_or_cross function to judge the accuracy.

@staticmethod
def tick_or_cross(predictions, labels):
    pred = torch.max(predictions.data, 1)[1]
    rights = pred.eq(labels.data.view_as(pred)).sum()
    return round(rights / len(pred), 4)
    
def evaluate(self, data_loader):
    accuracy = 0
    with torch.no_grad():
        for data in data_loader:
            d, label = data
            predictions = torch.argmax(self.net(d), dim=1)
            accuracy += self.tick_or_cross(predictions, label)
    print(f"accuracy: {
      
      accuracy / len(data_loader) * 100:.4f}%")

5 complete code

import torch
import torch.nn as nn
import torch.nn.functional as f
from torchvision.datasets import cifar
from torchvision import transforms
from torch.utils.data import DataLoader


class Net(nn.Module):
    """网络结构"""
    def __init__(self):
        """
        初始化网络结构和损失函数
        """
        super().__init__()
        # out_size = （in_size - k + 2p）/ s +1
        self.conv1 = nn.Conv2d(3, 32, (3, 3), (1, 1), 1)
        self.conv2 = nn.Conv2d(32, 64, (3, 3), (1, 1), 1)
        self.fc1 = nn.Linear(64 * 8 * 8, 64 * 8)
        self.fc2 = nn.Linear(64 * 8, 64)
        self.fc3 = nn.Linear(64, 10)

    def forward(self, x):
        """正向传播"""
        x = f.relu(f.max_pool2d(f.relu(self.conv1(x)), (2, 2)))
        x = f.relu(f.max_pool2d(f.relu(self.conv2(x)), (2, 2)))
        x = x.view(-1, 64 * 8 * 8)
        x = f.relu(self.fc1(x))
        x = f.relu(self.fc2(x))
        x = f.softmax(self.fc3(x))
        return x


class Model:
    """处理数据、训练、评估"""
    def __init__(self, net, learning_rate=0.001, batch_size=128, path=r'C:\Users\daifu\.keras\datasets'):
        """
        初始化网络结构、参数、数据存储路径
        :param net: 网络结构
        :param learning_rate: 学习率
        :param batch_size: 批量大小
        :param path: 数据存储路径
        """
        self.net = net
        self.cost = nn.CrossEntropyLoss()
        self.optimizer = torch.optim.RMSprop(self.net.parameters(), lr=learning_rate)
        self.batch_size = batch_size
        self.path = path

        # 数据增强
        self.transforms = transforms.Compose([
            transforms.CenterCrop([32, 32]),
            transforms.ToTensor(),
            transforms.Normalize((0.0, 0.5, 0.5), (0.5, 0.5, 0.5))
        ])

        # 获取数据
        self.train = cifar.CIFAR10(root=self.path, train=True, transform=self.transforms, download=True)
        self.train_loader = DataLoader(dataset=self.train, batch_size=self.batch_size, shuffle=True)
        self.test = cifar.CIFAR10(root=self.path, train=False, transform=self.transforms)
        self.test_loader = DataLoader(dataset=self.test, batch_size=self.batch_size, shuffle=True)

    @staticmethod
    def tick_or_cross(predictions, labels):
        pred = torch.max(predictions.data, 1)[1]
        rights = pred.eq(labels.data.view_as(pred)).sum()
        return round(rights / len(pred), 4)

    def train_(self, epochs=3):
        for epoch in range(epochs):
            losses = 0.0
            for i, data in enumerate(self.train_loader, 0):
                train_data, train_label = data
                self.optimizer.zero_grad()
                loss = self.cost(self.net(train_data), train_label)
                loss.backward()
                self.optimizer.step()
                losses += loss.item()
                if i % 10 == 0:
                    print(f'[epoch {
      
      epoch + 1}/{
      
      epochs}, {
      
      (i + 1) / len(self.train_loader) * 100:.2f}%] loss: {
      
      losses / 100:.2f}')
                    losses = 0.0
        print('Train Finished!')

    def evaluate(self, data_loader):
        accuracy = 0
        with torch.no_grad():
            for data in data_loader:
                d, label = data
                predictions = torch.argmax(self.net(d), dim=1)
                accuracy += self.tick_or_cross(predictions, label)
        print(f"accuracy: {
      
      accuracy / len(data_loader) * 100:.4f}%")


if __name__ == '__main__':
    _net = Net()
    model = Model(_net)
    model.train_(10)
    model.evaluate(model.test_loader)