Make writing a habit together! This is the 8th day of my participation in the “Gold Digging Day New Plan · April More Text Challenge”. Click here for more details.

@[toc]

Abstract

In this example, we extract some data from the seedling dataset as a data set. There are 12 categories in the data set. We demonstrate how to use pytorch version of MobileNetV1 image classification model to achieve classification tasks.

Through this article you will learn:

1, how to customize the MobileNetV1 model.

2. How to customize the loading method of data set?

3. How to use Cutout data enhancement?

4. How to use Mixup data enhancement.

5. How to implement training and verification.

6. Two ways of writing predictions.

MobileNets: An Efficient Convolutional Neural Network for Mobile Vision Applications

MobileNetV1 resolution:

MobileNetV1 network parsing, as well as implementation (PyTorch) _AIhao -CSDN blog

Keras version:

MobileNet combat: TensorFlow2.x version, MobileNetV1 image classification task (big data set) _AI Hao -CSDN blog

Data enhancement Cutout and Mixup

To improve my performance, I added Cutout and Mixup enhancements to the code. Implementing these two enhancements requires the installation of torchtoolbox. Installation command:

pip install torchtoolbox
Copy the code

The Cutout implementation, in transforms.

from torchtoolbox.transform import Cutout

# Data preprocessing

transform = transforms.Compose([
    transforms.Resize((224.224)),
    Cutout(),
    transforms.ToTensor(),
    transforms.Normalize([0.5.0.5.0.5], [0.5.0.5.0.5]])Copy the code

Mixup implementation, in the train method. From torchtoolbox.tools import mixup_data, mixup_criterion

    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
        data, labels_a, labels_b, lam = mixup_data(data, target, alpha)
        optimizer.zero_grad()
        output = model(data)
        loss = mixup_criterion(criterion, output, labels_a, labels_b, lam)
        loss.backward()
        optimizer.step()
        print_loss = loss.data.item()
Copy the code

The project structure

├─ Data │ ├─ Heavy Exercises │ ├─ Heavy Exercises │ ├─Common Chickweed │ ├─Common Wheat │ ├─ Heavy Exercises │ ├─ Heavy Exercises ├─Fat Hen │ ├─Loose Silky-bent │ ├─ ─Maize │ ├─Scentless Mayweed │ ├─Shepherds Purse │ ├─Small flowered Cranesbill │ └ ─ Sugar beet ├ ─ the dataset │ └ ─ dataset. Py └ ─ models │ └ ─ mobilenetV1. Py ├ ─ train. Py ├ ─ test1. Py └ ─ test. PyCopy the code

Import the libraries used by the project

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.utils.data
import torch.utils.data.distributed
import torchvision.transforms as transforms

from dataset.dataset import SeedlingData
from torch.autograd import Variable
from Model.mobilenetv1 import MobileNetV1
from torchtoolbox.tools import mixup_data, mixup_criterion
from torchtoolbox.transform import Cutout
Copy the code

Setting global Parameters

Set the learning rate, BatchSize, and epoch to check whether GPU exists in the environment. If not, use CPU. GPU is recommended, CPU is too slow.

Set global parameters
modellr = 1e-4
BATCH_SIZE = 16
EPOCHS = 300
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
Copy the code

Image preprocessing and enhancement

The data processing is relatively simple, including Cutout, Resize and normalization.

# Data preprocessing
transform = transforms.Compose([
    transforms.Resize((224.224)),
    Cutout(),
    transforms.ToTensor(),
    transforms.Normalize([0.5.0.5.0.5], [0.5.0.5.0.5])

])
transform_test = transforms.Compose([
    transforms.Resize((224.224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5.0.5.0.5], [0.5.0.5.0.5]])Copy the code

Read the data

Decompress the data set and put it under the data folder, as shown in the figure below:

Py and datasets.py will be added to the dataset folder and the following code will be added to the dataset folder:

# coding:utf8
import os
from PIL import Image
from torch.utils import data
from torchvision import transforms as T
from sklearn.model_selection import train_test_split
 
Labels = {'Black-grass': 0.'Charlock': 1.'Cleavers': 2.'Common Chickweed': 3.'Common wheat': 4.'Fat Hen': 5.'Loose Silky-bent': 6.'Maize': 7.'Scentless Mayweed': 8.'Shepherds Purse': 9.'Small-flowered Cranesbill': 10.'Sugar beet': 11}
 
 
class SeedlingData (data.Dataset) :
 
    def __init__(self, root, transforms=None, train=True, test=False) :
        "" Main objective: to obtain the address of all images and divide the data according to training, validation and testing.
        self.test = test
        self.transforms = transforms
 
        if self.test:
            imgs = [os.path.join(root, img) for img in os.listdir(root)]
            self.imgs = imgs
        else:
            imgs_labels = [os.path.join(root, img) for img in os.listdir(root)]
            imgs = []
            for imglable in imgs_labels:
                for imgname in os.listdir(imglable):
                    imgpath = os.path.join(imglable, imgname)
                    imgs.append(imgpath)
            trainval_files, val_files = train_test_split(imgs, test_size=0.3, random_state=42)
            if train:
                self.imgs = trainval_files
            else:
                self.imgs = val_files
 
    def __getitem__(self, index) :
        """ Return data one image at a time ""
        img_path = self.imgs[index]
        img_path=img_path.replace("\ \".'/')
        if self.test:
            label = -1
        else:
            labelname = img_path.split('/')[-2]
            label = Labels[labelname]
        data = Image.open(img_path).convert('RGB')
        data = self.transforms(data)
        return data, label
 
    def __len__(self) :
        return len(self.imgs)
Copy the code

The core logic of the code:

The first step is to create a dictionary and define the ID of the category, replacing the category with a number.

The second step is to write a method to get the image path in __init__. There is only one layer path for test set to read directly. The training set is the category folder under the train folder. The category is obtained first, and then the specific picture path is obtained. The training set and verification set are then split in a 7:3 ratio using the sklearn method of dividing the data set.

The third step is to define the method to read individual images and categories in the __getitem__ method. Since there is a 32-bit depth in the image, I made the conversion when reading the image.

Then we call SeedlingData on train.py to read the data. Remember to import the dataset. Py (from dataset.

dataset_train = SeedlingData('data/train', transforms=transform, train=True)
dataset_test = SeedlingData("data/train", transforms=transform_test, train=False)
# fetch data
print(dataset_train.imgs)

# import data
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=BATCH_SIZE, shuffle=False)

Copy the code

Set up the model

• Set the loss function to nn.CrossentRopyLoss ().
• Set the model to MobileNetV1 and num_classes to 12.
• The optimizer is set to Adam.
• The learning rate adjustment strategy is cosine annealing.

Instantiate the model and move it to the GPU
criterion = nn.CrossEntropyLoss()
model_ft = MobileNetV1(num_classes=12)
model_ft.to(DEVICE)
# Choose simple violent Adam optimizer, learning rate down
optimizer = optim.Adam(model_ft.parameters(), lr=modellr)
cosine_schedule = optim.lr_scheduler.CosineAnnealingLR(optimizer=optimizer,T_max=20,eta_min=1e-9)
Copy the code

Define training and validation functions

# Define the training process
alpha=0.2
def train(model, device, train_loader, optimizer, epoch) :
    model.train()
    sum_loss = 0
    total_num = len(train_loader.dataset)
    print(total_num, len(train_loader))
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device, non_blocking=True), target.to(device, non_blocking=True)
        data, labels_a, labels_b, lam = mixup_data(data, target, alpha)
        optimizer.zero_grad()
        output = model(data)
        loss = mixup_criterion(criterion, output, labels_a, labels_b, lam)
        loss.backward()
        optimizer.step()
        lr = optimizer.state_dict()['param_groups'] [0] ['lr']
        print_loss = loss.data.item()
        sum_loss += print_loss
        if (batch_idx + 1) % 10= =0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLR:{:.9f}'.format(
                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                       100. * (batch_idx + 1) / len(train_loader), loss.item(),lr))
    ave_loss = sum_loss / len(train_loader)
    print('epoch:{},loss:{}'.format(epoch, ave_loss))

ACC=0
# Validation process
def val(model, device, test_loader) :
    global ACC
    model.eval()
    test_loss = 0
    correct = 0
    total_num = len(test_loader.dataset)
    print(total_num, len(test_loader))
    with torch.no_grad():
        for data, target in test_loader:
            data, target = Variable(data).to(device), Variable(target).to(device)
            output = model(data)
            loss = criterion(output, target)
            _, pred = torch.max(output.data, 1)
            correct += torch.sum(pred == target)
            print_loss = loss.data.item()
            test_loss += print_loss
        correct = correct.data.item()
        acc = correct / total_num
        avgloss = test_loss / len(test_loader)
        print('\nVal set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            avgloss, correct, len(test_loader.dataset), 100 * acc))
        if acc > ACC:
            torch.save(model_ft, 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')
            ACC = acc


# training

for epoch in range(1, EPOCHS + 1):
    train(model_ft, DEVICE, train_loader, optimizer, epoch)
    cosine_schedule.step()
    val(model_ft, DEVICE, test_loader)

Copy the code

Running results:

test

I will introduce two common testing methods. The first is universal, which is to manually load data sets and make predictions. The specific operations are as follows:

The test set is stored in the following directory:

The first step is to define the category, the order of this category and the training of the category order corresponding, do not change the order ！！！！

Second, define transforms. Transforms are the same as the validation set’s transforms, without data enhancement.

Step 3 load the model and put it in DEVICE,

The fourth step is to read the Image and predict the category of the Image. Note here that Image is read using PIL library Image. Don’t use CV2, transforms is not supported.

import torch.utils.data.distributed
import torchvision.transforms as transforms
from PIL import Image
from torch.autograd import Variable
import os
classes = ('Black-grass'.'Charlock'.'Cleavers'.'Common Chickweed'.'Common wheat'.'Fat Hen'.'Loose Silky-bent'.'Maize'.'Scentless Mayweed'.'Shepherds Purse'.'Small-flowered Cranesbill'.'Sugar beet')
transform_test = transforms.Compose([
         transforms.Resize((224.224)),
        transforms.ToTensor(),
        transforms.Normalize([0.5.0.5.0.5], [0.5.0.5.0.5])
])
 
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)
 
path='data/test/'
testList=os.listdir(path)
for file in testList:
        img=Image.open(path+file)
        img=transform_test(img)
        img.unsqueeze_(0)
        img = Variable(img).to(DEVICE)
        out=model(img)
        # Predict
        _, pred = torch.max(out.data, 1)
        print('Image Name:{},predict:{}'.format(file,classes[pred.data.item()]))
Copy the code

Running results:

The second uses a custom Dataset to read the image

import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import SeedlingData
from torch.autograd import Variable
 
classes = ('Black-grass'.'Charlock'.'Cleavers'.'Common Chickweed'.'Common wheat'.'Fat Hen'.'Loose Silky-bent'.'Maize'.'Scentless Mayweed'.'Shepherds Purse'.'Small-flowered Cranesbill'.'Sugar beet')
transform_test = transforms.Compose([
    transforms.Resize((224.224)),
    transforms.ToTensor(),
    transforms.Normalize([0.5.0.5.0.5], [0.5.0.5.0.5])
])
 
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = torch.load("model.pth")
model.eval()
model.to(DEVICE)
 
dataset_test =SeedlingData('data/test/', transform_test,test=True)
print(len(dataset_test))
# Label of the corresponding folder
 
for index in range(len(dataset_test)):
    item = dataset_test[index]
    img, label = item
    img.unsqueeze_(0)
    data = Variable(img).to(DEVICE)
    output = model(data)
    _, pred = torch.max(output.data, 1)
    print('Image Name:{},predict:{}'.format(dataset_test.imgs[index], classes[pred.data.item()]))
    index += 1
 
Copy the code

Running results:

Complete code:Download.csdn.net/download/hh…