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 codeSetting 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 codeImage 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 codeRead 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 codeThen 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 codeSet 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 codeSet 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 codetest
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 codeThe 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 codeThe 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 codeVGG implements plant seedling classification