(1) Pytorch single task image classification

Deep learning mainly consists of four parts: data reading, network model, loss function, and optimizer.

You shouldn’t get hung up on these details at the beginning. You should let the code run first and then study how the code is written.

The following code is only the training part of the code, plus the function of verifying the accuracy of the model.

1. Project distribution: Create a folder my_data1, and create the train and valid folders in my_data1

(The folder name is fixed, don’t write train and valid incorrectly, otherwise the code will not run) .

  train is the picture of the training set, and valid is the picture of the verification set. In the train folder, create as many folders as you need for training ( for example, if I only train two categories, I will only create two folders, cat and dog, and the folder names are not fixed )

  The format of the valid folder is the same as that of train.

2. Code parameter introduction: If the category you train is 3, change num_classes=2 in the code to num_classes=3

   The code only trains for 100 rounds by default. If you want to train for 200 rounds, change Epoches=100 in the code to Epoches=200.

   The default Batch_size of the code is 4. The setting depends on the graphics card. The better the graphics card, the larger the value you can set.

   By default, the code resizes the image to [224, 224] before training. If you want to change it, you can modify Image_Size.

3. Training code train.py: The model used in it is resnet18, and the pre-trained model is loaded for training, and then the first 30 layers are frozen.

import torch
from torch.autograd import Variable
import torchvision
from torchvision import datasets, transforms, models
import os
import matplotlib.pyplot as plt
import time
import torch.optim as optim

from PIL import Image, ImageFile

ImageFile.LOAD_TRUNCATED_IMAGES = True
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def train():
    running_loss = 0
    for batch_idx, (data, target) in enumerate(train_data):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        out = net(data)
        loss = criterion(out, target)
        running_loss += loss.item()
        loss.backward()
        optimizer.step()

    return running_loss


def test():
    correct, total = 0, 0
    with torch.no_grad():
        for _, (data, target) in enumerate(val_data):
            data, target = data.to(device), target.to(device)
            out = net(data)
            prediction = torch.max(out.data, dim=1)[1]
            total += target.size(0)
            correct += (prediction == target).sum().item()
        print('Accuracy on test set: (%d/%d)=%d %%' % (correct, total, 100 * correct / total))


if __name__ == '__main__':
    loss_list = []
    Epoches = 100
    Batch_Size = 4
    Image_Size = [224, 224]

    # 1.数据加载
    data_dir = r'D:\Code\python\完整项目放置\classify_project\multi_classification\my_dataset1'
    # 1.1 定义要对数据进行的处理
    data_transform = {x: transforms.Compose([transforms.Resize(Image_Size), transforms.ToTensor()]) for x in
                      ["train", "valid"]}
    image_datasets = {x: datasets.ImageFolder(root=os.path.join(data_dir, x), transform=data_transform[x]) for x in
                      ["train", "valid"]}
    dataloader = {x: torch.utils.data.DataLoader(dataset=image_datasets[x], batch_size=Batch_Size, shuffle=True) for x in
                  ["train", "valid"]}
    train_data, val_data = dataloader["train"], dataloader["valid"]

    index_classes = image_datasets["train"].class_to_idx
    print(index_classes)
    example_classes = image_datasets["train"].classes
    print(example_classes)

    # 2.数据预览, 在训练的时候可以注释掉
    # X_example, y_example = next(iter(dataloader["train"]))
    # img = torchvision.utils.make_grid(X_example)
    # img = img.numpy().transpose([1, 2, 0])
    # for i in range(len(y_example)):
    #     index = y_example[i]
    #     print(example_classes[index], end='   ')
    #     if (i+1)%8 == 0:
    #         print()
    # plt.imshow(img)
    # plt.show()

    # 3.模型加载, 并对模型进行微调
    net = models.resnet18(pretrained=True)
    fc_features = net.fc.in_features

    # 设置训练的类别个数，我这里只有两类所以写2
    num_classes = 2
    net.fc = torch.nn.Linear(fc_features, num_classes)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 4.pytorch fine tune 微调(冻结一部分层)。这里是冻结网络前30层参数进行训练。
    for i, param in enumerate(net.parameters()):
        if i < 30:
            param.requires_grad = False
    net.to(device)

    # 5.定义损失函数，以及优化器
    LR = 1e-3
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=LR)

    for epoch in range(Epoches):
        loss = train()
        loss_list.append(loss)
        print("第%d轮的loss为：%5f:" % (epoch, loss))
        test()

        # net.state_dict只保存模型的参数
        # torch.save(net.state_dict(), 'Model2.pth')

        # 保存整个模型
        torch.save(net, "my_model.pth")

    plt.title("Graph")
    plt.plot(range(Epoches), loss_list)
    plt.ylabel("loss")
    plt.xlabel("epoch")
    plt.show()