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

test

Complete code:

---

Abstract

EfficientNet EfficientNet is Google’s 2019 classification model, and since then it has been one of the most EfficientNet EfficientNet EfficientNet platforms. The following figure EfficientNet – B0 model network structure.

As can be seen from the network, the author constructed MBConv with the structure as shown in the figure below:

The size of the convolution kernel corresponding to K is amplified by 4 times after 1×1 convolution, and then through depthwise conv3×3 convolution, and then through SE module, and then through 1×1 convolution, the channel is restored to the size of the input, and finally fused with the upper input.

EfficientNet EfficientNet EfficientNet EfficientNet EfficientNet EfficientNet EfficientNet EfficientNet

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

Check whether the EfficientNet library has been installed. If not, run PIP Install Efficientnet_Pytorch to install and import the EfficientNet library.

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
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 efficientnet_pytorch import EfficientNet
#pip install efficientnet_pytorch
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.

Modellr = 1E-4 BATCH_SIZE = 64 EPOCHS = 20 DEVICE = torch. DEVICE (‘cuda’ if torch. Cuda.is_available () else ‘cpu’)

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.

Transforms = 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. 224)), transforms. ToTensor (), transforms. The Normalize ([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])

Read the data

Address: data sets link: pan.baidu.com/s/1ZM8vDWEz… Download it and unzip it into the data folder. The data catalog is shown below:

​

Py = __init__. Py = dataset. Py = __init__. Py = dataset.

# 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

class DogCat(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
        imgs = [os.path.join(root, img) for img in os.listdir(root)]
        if self.test:
            imgs = sorted(imgs, key=lambda x: int(x.split('. ')[-2].split('/')[-1]))
        else:
            imgs = sorted(imgs, key=lambda x: int(x.split('. ')[-2]))

        if self.test:
            self.imgs = imgs
        else:
            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]
        if self.test:
            label =-1
        else:
            label = 1 if 'dog' in img_path.split('/')[-1] else 0
        data = Image.open(img_path)
        data = self.transforms(data)
        return data, label

    def __len__(self) :
        return len(self.imgs)
Copy the code

We then call DogCat on train.py to read the data

dataset_train = DogCat('data/train', transforms=transform, train=True)
dataset_test = DogCat("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 and the model is efficientnet-B3. Change the last layer’s full connection, set the category to 2, and place the model on DEVICE. The optimizer selects Adam.

Instantiate the model and move it to the GPU
criterion = nn.CrossEntropyLoss()
model_ft = EfficientNet.from_pretrained('efficientnet-b3')
num_ftrs = model_ft._fc.in_features
model_ft._fc = nn.Linear(num_ftrs, 2)
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) % 50= =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

After completing the above code, you can start training. Click Run to start training, as shown below:

Because we use the pre-training model, the convergence rate is very fast.

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 ！！！！ When we train, cat is 0,dog is 1, so I define classes as (cat,dog).

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
import torchvision.datasets as datasets
from PIL import Image
from torch.autograd import Variable
import os
classes = ('cat'.'dog')
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 the dataset. Py we just defined to load the test set. The code is as follows:

import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import DogCat
from torch.autograd import Variable

classes = ('cat'.'dog')
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 =DogCat('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

Complete code:

train.py

import torch.optim as optim
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.optim
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 efficientnet_pytorch import EfficientNet
#pip install efficientnet_pytorch

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 = DogCat('data/train', transforms=transform, train=True)
dataset_test = DogCat("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 = EfficientNet.from_pretrained('efficientnet-b3')
num_ftrs = model_ft._fc.in_features
model_ft._fc = nn.Linear(num_ftrs, 2)
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) % 50= =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

test1.py

import torch.utils.data.distributed
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from PIL import Image
from torch.autograd import Variable
import os
classes = ('cat'.'dog')
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

test2.py

import torch.utils.data.distributed
import torchvision.transforms as transforms
from dataset.dataset import DogCat
from torch.autograd import Variable

classes = ('cat'.'dog')
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 =DogCat('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

Image classification EfficientNet EfficientNet EfficientNet