Hoy se realizó un trabajo de clasificación de imágenes utilizando resnet50, y el código ya está grabado. En la primera mitad de este artículo, construya rápidamente una red resnet para entrenamiento. Esta parte del código se refiere principalmente al código completo del blog usando la red resnet50 para entrenar modelos de clasificación múltiple . He hecho ciertas modificaciones y agregado comentarios a el código basado en este blog; además, en este blog La última parte del contenido es principalmente un script escrito por mí mismo para construir rápidamente el modelo de red resnet50 previamente entrenado y realizar la detección de imágenes.
1. El modelo resnet50 se construye rápidamente y el conjunto de datos se entrena
La implementación rápida aquí utiliza principalmente la red preentrenada resnet50 que viene con pytorch, lo que ahorra el proceso de agregar estructuras a su vez.
Con respecto a la estructura de almacenamiento de carpetas del conjunto de datos antes de la preparación del entrenamiento, como se describe en el blog original:
--tren
|--perro
|--gato...
--válido
|--perro
|--gato
Cree una carpeta de datos (el nombre puede ser arbitrario) en el directorio raíz del proyecto para almacenar el conjunto de datos.
Cree carpetas de tren, válidas y de prueba en secuencia debajo de la carpeta de datos (la carpeta de prueba puede estar ausente, determinada según sus propias necesidades)
Cree carpetas de categoría debajo de la carpeta de tren, como gato, perro, etc.
, y guárdelas en carpetas de categoría como gato Una imagen de la categoría gato.
…
Cree carpetas de categoría en la carpeta val, como gato, perro, etc.
En la carpeta de categoría, como gato, almacene imágenes de la categoría gato.
…
——————————————————
Nota
: Esta cita se extrajo en parte del artículo original del bloguero de CSDN "Trabajando en silencio"
Enlace del artículo original: https://blog.csdn.net /weixin_46034990 /artículo/detalles/124859877
Además, se debe tener en cuenta que se debe crear un nuevo archivo llamado modelos en el directorio raíz de la carpeta de código para guardar el archivo del modelo durante el proceso de capacitación. Al mismo tiempo, según su propio conjunto de entrenamiento, debe personalizar y modificar la variable num_classes del número de categorías de la clasificación de destino en el código.
A continuación se muestra el código de entrenamiento e implementación del modelo:
import torch
from torchvision import datasets, models, transforms
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import time
import numpy as np
import matplotlib.pyplot as plt
import os
from tqdm import tqdm
# 一、建立数据集
# animals-6
# --train
# |--dog
# |--cat
# ...
# --valid
# |--dog
# |--cat
# ...
# --test
# |--dog
# |--cat
# ...
# 我的数据集中 train 中每个类别60张图片，valid 中每个类别 10 张图片，test 中每个类别几张到几十张不等，一共 6 个类别。
# 二、数据增强
# 建好的数据集在输入网络之前先进行数据增强，包括随机 resize 裁剪到 256 x 256，随机旋转，随机水平翻转，中心裁剪到 224 x 224，转化成 Tensor，正规化等。
image_transforms = {
'train': transforms.Compose([
transforms.RandomResizedCrop(size=256, scale=(0.8, 1.0)),
transforms.RandomRotation(degrees=15),
transforms.RandomHorizontalFlip(),
transforms.CenterCrop(size=224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
]),
'valid': transforms.Compose([
transforms.Resize(size=256),
transforms.CenterCrop(size=224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
}
# 三、加载数据
# torchvision.transforms包DataLoader是 Pytorch 重要的特性，它们使得数据增加和加载数据变得非常简单。
# 使用 DataLoader 加载数据的时候就会将之前定义的数据 transform 就会应用的数据上了。
dataset = 'data'
train_directory = os.path.join(dataset, 'train')
valid_directory = os.path.join(dataset, 'valid')
batch_size = 32
num_classes = 2 #分类种类数
print(train_directory)
data = {
'train': datasets.ImageFolder(root=train_directory, transform=image_transforms['train']),
'valid': datasets.ImageFolder(root=valid_directory, transform=image_transforms['valid'])
}
print("训练集图片类别及其对应编号(种类名：编号)：",data['train'].class_to_idx)
print("测试集图片类别及其对应编号：",data['valid'].class_to_idx)
train_data_size = len(data['train'])
valid_data_size = len(data['valid'])
train_data = DataLoader(data['train'], batch_size=batch_size, shuffle=True, num_workers=8)
valid_data = DataLoader(data['valid'], batch_size=batch_size, shuffle=True, num_workers=8)
print("训练集图片数量：",train_data_size, "测试集图片数量：",valid_data_size)
# 四、迁移学习
# 这里使用ResNet-50的预训练模型。
resnet50 = models.resnet50(pretrained=True)
# 在PyTorch中加载模型时，所有参数的‘requires_grad’字段默认设置为true。这意味着对参数值的每一次更改都将被存储，以便在用于训练的反向传播图中使用。
# 这增加了内存需求。由于预训练的模型中的大多数参数已经训练好了，因此将requires_grad字段重置为false。
for param in resnet50.parameters():
param.requires_grad = False
# 为了适应自己的数据集，将ResNet-50的最后一层替换为，将原来最后一个全连接层的输入喂给一个有256个输出单元的线性层，接着再连接ReLU层和Dropout层，然后是256 x 6的线性层，输出为6通道的softmax层。
fc_inputs = resnet50.fc.in_features
resnet50.fc = nn.Sequential(
nn.Linear(fc_inputs, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, num_classes),
nn.LogSoftmax(dim=1)
)
# 用GPU进行训练。
resnet50 = resnet50.to('cuda:0')
# 定义损失函数和优化器。
loss_func = nn.NLLLoss()
optimizer = optim.Adam(resnet50.parameters())
# 五、训练
def train_and_valid(model, loss_function, optimizer, epochs=25):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
history = []
best_acc = 0.0
best_epoch = 0
for epoch in range(epochs):
epoch_start = time.time()
print("Epoch: {}/{}".format(epoch+1, epochs))
model.train()
train_loss = 0.0
train_acc = 0.0
valid_loss = 0.0
valid_acc = 0.0
for i, (inputs, labels) in enumerate(tqdm(train_data)):
inputs = inputs.to(device)
labels = labels.to(device)
#因为这里梯度是累加的，所以每次记得清零
optimizer.zero_grad()
outputs = model(inputs)
loss = loss_function(outputs, labels)
print("标签值：",labels)
print("输出值：",outputs)
loss.backward()
optimizer.step()
train_loss += loss.item() * inputs.size(0)
ret, predictions = torch.max(outputs.data, 1)
correct_counts = predictions.eq(labels.data.view_as(predictions))
acc = torch.mean(correct_counts.type(torch.FloatTensor))
train_acc += acc.item() * inputs.size(0)
with torch.no_grad():
model.eval()
for j, (inputs, labels) in enumerate(tqdm(valid_data)):
inputs = inputs.to(device)
labels = labels.to(device)
outputs = model(inputs)
loss = loss_function(outputs, labels)
valid_loss += loss.item() * inputs.size(0)
ret, predictions = torch.max(outputs.data, 1)
correct_counts = predictions.eq(labels.data.view_as(predictions))
acc = torch.mean(correct_counts.type(torch.FloatTensor))
valid_acc += acc.item() * inputs.size(0)
avg_train_loss = train_loss/train_data_size
avg_train_acc = train_acc/train_data_size
avg_valid_loss = valid_loss/valid_data_size
avg_valid_acc = valid_acc/valid_data_size
history.append([avg_train_loss, avg_valid_loss, avg_train_acc, avg_valid_acc])
if best_acc < avg_valid_acc:
best_acc = avg_valid_acc
best_epoch = epoch + 1
epoch_end = time.time()
print("Epoch: {:03d}, Training: Loss: {:.4f}, Accuracy: {:.4f}%, \n\t\tValidation: Loss: {:.4f}, Accuracy: {:.4f}%, Time: {:.4f}s".format(
epoch+1, avg_valid_loss, avg_train_acc*100, avg_valid_loss, avg_valid_acc*100, epoch_end-epoch_start
))
print("Best Accuracy for validation : {:.4f} at epoch {:03d}".format(best_acc, best_epoch))
torch.save(model, 'models/'+dataset+'_model_'+str(epoch+1)+'.pt')
return model, history
num_epochs = 3 #训练周期数
trained_model, history = train_and_valid(resnet50, loss_func, optimizer, num_epochs)
torch.save(history, 'models/'+dataset+'_history.pt')
history = np.array(history)
plt.plot(history[:, 0:2])
plt.legend(['Tr Loss', 'Val Loss'])
plt.xlabel('Epoch Number')
plt.ylabel('Loss')
plt.ylim(0, 1)
plt.savefig(dataset+'_loss_curve.png')
plt.show()
plt.plot(history[:, 2:4])
plt.legend(['Tr Accuracy', 'Val Accuracy'])
plt.xlabel('Epoch Number')
plt.ylabel('Accuracy')
plt.ylim(0, 1)
plt.savefig(dataset+'_accuracy_curve.png')
plt.show()
2. Vuelva a extraer el archivo guardado del modelo entrenado cuando lo use, y transfiera la imagen para ser detectada al modelo entrenado y lea el resultado.
Esta parte es principalmente para leer el archivo guardado del modelo entrenado en la parte anterior (este paso necesita construir un modelo con la misma estructura que el modelo que usamos para el entrenamiento e importar los pesos del modelo entrenado a este nuevo modelo), y luego ingrese la imagen que se detectará en el modelo para la detección, e interprete y lea el resultado de salida del modelo (el resultado de salida es un vector de tensor).
El código de proceso específico es el siguiente:
import torch
from torchvision import models, transforms
import torch.nn as nn
import cv2
classes = ["1","2"] #识别种类名称（顺序要与训练时的数据导入编号顺序对应，可以使用datasets.ImageFolder().class_to_idx来查看）
transf = transforms.ToTensor()
device = torch.device('cuda:0')
num_classes = 2
model_path = "models/data_model_3.pt"
image_input = cv2.imread("test_image1.jpg")
image_input = transf(image_input)
image_input = torch.unsqueeze(image_input,dim=0).cuda()
#搭建模型
resnet50 = models.resnet50(pretrained=True)
for param in resnet50.parameters():
param.requires_grad = False
fc_inputs = resnet50.fc.in_features
resnet50.fc = nn.Sequential(
nn.Linear(fc_inputs, 256),
nn.ReLU(),
nn.Dropout(0.4),
nn.Linear(256, num_classes),
nn.LogSoftmax(dim=1)
)
resnet50 = torch.load(model_path)
outputs = resnet50(image_input)
value,id =torch.max(outputs,1)
print(outputs,"\n","结果是：",classes[id])