directory

Abstract

New project

Import the required libraries

Setting global Parameters

Image preprocessing

Read the data

Set up the model

Set up training and validation

The complete code

Abstract

VGG16 this time we use the classic image classification model, the implementation of the classification of plant seedlings, the data collection of links: pan.baidu.com/s/1JIczDc7V… , there are 12 categories. Here is a sample image.

Most of the images are 24 bit deep, and some are 32 bit, so you have to cast the images. Here is a point to remind you, get the data set, do not come up with algorithms, first to browse the data set, understand what the data set looks like, how many pictures, how easy to identify a preliminary understanding.

The model adopts VGG, and the detailed introduction of the model refer to the article VGGNet (PyTorch) _AIhao -CSDN blog.

Next, how to use VGG to classify plant seedlings.

New project

Create a project for image classification, put data sets in data, and customize the data reading method in dataset folder. I don’t use the default reading method this time, because it is too simple and meaningless. Then create new train.py and test.py

Create a new train. Py at the root of the project and write the training code inside.

Import the required libraries

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 DogCat
from torch.autograd import Variable
from torchvision.models import vgg16
Copy the code

Setting global Parameters

Set BatchSize, learning rate, and EPOchs to check whether the CUDA environment exists. If not, set the BatchSize to CPU.

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

Image preprocessing

In image and processing, the transform of train data set and the transform of verification set are separated. Besides resize and normalization of train image processing, image enhancement can be set, such as rotation, random erasure and a series of operations, while image enhancement is not required for verification set. In addition, do not blindly enhance, unreasonable enhancement means are likely to bring negative effects, and even Loss does not converge

# Data preprocessing

transform = transforms.Compose([
    transforms.Resize((224.224)),
    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 = __init__. Py = dataset. Py = __init__. Py = dataset.

Talk about 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.

# 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

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

CrossEntropyLoss is used as loss, Alexnet is used as model, and pre-training model is selected. Change the full connection layer, set the last layer category to 12, and place the model on DEVICE. The optimizer selects Adam.

Instantiate the model and move it to the GPU
criterion = nn.CrossEntropyLoss()
model_ft = vgg16(pretrained=True)
model_ft.classifier = classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7.4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096.4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096.12),
        )
model_ft.to(DEVICE)
# Choose simple violent Adam optimizer, learning rate down
optimizer = optim.Adam(model_ft.parameters(), lr=modellr)


def adjust_learning_rate(optimizer, epoch) :
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    modellrnew = modellr * (0.1 ** (epoch // 50))
    print("lr:", modellrnew)
    for param_group in optimizer.param_groups:
        param_group['lr'] = modellrnew
Copy the code

Set up training and validation

# Define the training process

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 = Variable(data).to(device), Variable(target).to(device)
        output = model(data)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print_loss = loss.data.item()
        sum_loss += print_loss
        if (batch_idx + 1) % 10= =0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                       100. * (batch_idx + 1) / len(train_loader), loss.item()))
    ave_loss = sum_loss / len(train_loader)
    print('epoch:{},loss:{}'.format(epoch, ave_loss))


# Validation process
def val(model, device, test_loader) :
    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))


# training

for epoch in range(1, EPOCHS + 1) : adjust_learning_rate(optimizer, epoch) train(model_ft, DEVICE, train_loader, optimizer, epoch) val(model_ft, DEVICE, test_loader) torch.save(model_ft,'model.pth')
Copy the code

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

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

The complete code

train.py

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 torchvision.models import vgg16

Set global parameters
modellr = 1e-4
BATCH_SIZE = 32
EPOCHS = 10
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# Data preprocessing

transform = transforms.Compose([
    transforms.Resize((224.224)),
    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])
])
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)

Instantiate the model and move it to the GPU
criterion = nn.CrossEntropyLoss()
model_ft = vgg16(pretrained=True)
model_ft.classifier = classifier = nn.Sequential(
            nn.Linear(512 * 7 * 7.4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096.4096),
            nn.ReLU(True),
            nn.Dropout(),
            nn.Linear(4096.12),
        )
model_ft.to(DEVICE)
# Choose simple violent Adam optimizer, learning rate down
optimizer = optim.Adam(model_ft.parameters(), lr=modellr)


def adjust_learning_rate(optimizer, epoch) :
    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
    modellrnew = modellr * (0.1 ** (epoch // 50))
    print("lr:", modellrnew)
    for param_group in optimizer.param_groups:
        param_group['lr'] = modellrnew


# Define the training process

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 = Variable(data).to(device), Variable(target).to(device)
        output = model(data)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        print_loss = loss.data.item()
        sum_loss += print_loss
        if (batch_idx + 1) % 10= =0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, (batch_idx + 1) * len(data), len(train_loader.dataset),
                       100. * (batch_idx + 1) / len(train_loader), loss.item()))
    ave_loss = sum_loss / len(train_loader)
    print('epoch:{},loss:{}'.format(epoch, ave_loss))


# Validation process
def val(model, device, test_loader) :
    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))


# training

for epoch in range(1, EPOCHS + 1) : adjust_learning_rate(optimizer, epoch) train(model_ft, DEVICE, train_loader, optimizer, epoch) val(model_ft, DEVICE, test_loader) torch.save(model_ft,'model.pth')
Copy the code

VGG implements plant seedling classification