El aprendizaje profundo consta principalmente de cuatro partes: lectura de datos, modelo de red, función de pérdida y optimizador.
No debes obsesionarte con estos detalles al principio. Primero debes dejar que el código se ejecute y luego estudiar cómo está escrito.
El siguiente código es solo la parte de entrenamiento del código, más la función de verificar la precisión del modelo.
1. Distribución del proyecto: cree una carpeta my_data1 y cree el tren y las carpetas válidas en my_data1
(El nombre de la carpeta es fijo, no escriba train y valide incorrectamente; de lo contrario, el código no se ejecutará) .
train es la imagen del conjunto de entrenamiento y valid es la imagen del conjunto de verificación. En la carpeta de entrenamiento, cree tantas carpetas como necesite para entrenar ( por ejemplo, si solo entreno dos categorías, solo crearé dos carpetas, gato y perro, y los nombres de las carpetas no son fijos ).
El formato de la carpeta válida es el mismo que el de tren.
2. Introducción del parámetro de código: si la categoría que entrena es 3, cambie num_classes=2 en el código a num_classes=3
El código solo entrena para 100 rondas de forma predeterminada. Si desea entrenar para 200 rondas, cambie Epoches=100 en el código a Epoches=200.
El tamaño de lote predeterminado del código es 4. La configuración depende de la tarjeta gráfica. Cuanto mejor sea la tarjeta gráfica, mayor será el valor que podrá establecer.
De forma predeterminada, el código cambia el tamaño de la imagen a [224, 224] antes del entrenamiento. Si desea cambiarlo, puede modificar Image_Size.
3. Código de entrenamiento train.py: el modelo utilizado es resnet18, el modelo previamente entrenado se carga para el entrenamiento y luego se congelan las primeras 30 capas.
import torch
from torch.autograd import Variable
import torchvision
from torchvision import datasets, transforms, models
import os
import matplotlib.pyplot as plt
import time
import torch.optim as optim
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
def train():
running_loss = 0
for batch_idx, (data, target) in enumerate(train_data):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
out = net(data)
loss = criterion(out, target)
running_loss += loss.item()
loss.backward()
optimizer.step()
return running_loss
def test():
correct, total = 0, 0
with torch.no_grad():
for _, (data, target) in enumerate(val_data):
data, target = data.to(device), target.to(device)
out = net(data)
prediction = torch.max(out.data, dim=1)[1]
total += target.size(0)
correct += (prediction == target).sum().item()
print('Accuracy on test set: (%d/%d)=%d %%' % (correct, total, 100 * correct / total))
if __name__ == '__main__':
loss_list = []
Epoches = 100
Batch_Size = 4
Image_Size = [224, 224]
# 1.数据加载
data_dir = r'D:\Code\python\完整项目放置\classify_project\multi_classification\my_dataset1'
# 1.1 定义要对数据进行的处理
data_transform = {x: transforms.Compose([transforms.Resize(Image_Size), transforms.ToTensor()]) for x in
["train", "valid"]}
image_datasets = {x: datasets.ImageFolder(root=os.path.join(data_dir, x), transform=data_transform[x]) for x in
["train", "valid"]}
dataloader = {x: torch.utils.data.DataLoader(dataset=image_datasets[x], batch_size=Batch_Size, shuffle=True) for x in
["train", "valid"]}
train_data, val_data = dataloader["train"], dataloader["valid"]
index_classes = image_datasets["train"].class_to_idx
print(index_classes)
example_classes = image_datasets["train"].classes
print(example_classes)
# 2.数据预览, 在训练的时候可以注释掉
# X_example, y_example = next(iter(dataloader["train"]))
# img = torchvision.utils.make_grid(X_example)
# img = img.numpy().transpose([1, 2, 0])
# for i in range(len(y_example)):
# index = y_example[i]
# print(example_classes[index], end=' ')
# if (i+1)%8 == 0:
# print()
# plt.imshow(img)
# plt.show()
# 3.模型加载, 并对模型进行微调
net = models.resnet18(pretrained=True)
fc_features = net.fc.in_features
# 设置训练的类别个数，我这里只有两类所以写2
num_classes = 2
net.fc = torch.nn.Linear(fc_features, num_classes)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 4.pytorch fine tune 微调(冻结一部分层)。这里是冻结网络前30层参数进行训练。
for i, param in enumerate(net.parameters()):
if i < 30:
param.requires_grad = False
net.to(device)
# 5.定义损失函数，以及优化器
LR = 1e-3
criterion = torch.nn.CrossEntropyLoss()
optimizer = optim.Adam(filter(lambda p: p.requires_grad, net.parameters()), lr=LR)
for epoch in range(Epoches):
loss = train()
loss_list.append(loss)
print("第%d轮的loss为：%5f:" % (epoch, loss))
test()
# net.state_dict只保存模型的参数
# torch.save(net.state_dict(), 'Model2.pth')
# 保存整个模型
torch.save(net, "my_model.pth")
plt.title("Graph")
plt.plot(range(Epoches), loss_list)
plt.ylabel("loss")
plt.xlabel("epoch")
plt.show()