【AI实战】Pytorch版使用预训练卷积神经网络快速训练自己的分类模型
Pytorch版使用预训练模型快速训练自己的分类模型
-
- 常见的卷积神经网络
- 预训练模型
- 使用预训练的 ResNet 18 快速训练自己的分类模型核心步骤
- 使用预训练的 ResNet 18 快速训练自己的分类模型详细过程
- 完整代码
常见的卷积神经网络
- alexnet
- vgg
- resnet
- inception
- densenet
- googlenet
预训练模型
ResNet 有很多变种，包括 ResNet 18、ResNet 34、ResNet 50、ResNet 101、ResNet 152。
这里以 ResNet 18 为例子详细说明Pytorch版使用预训练模型快速训练自己的分类模型的过程。
使用预训练的 ResNet 18 快速训练自己的分类模型核心步骤
- 1、首先加载训练好的模型参数：
resnet18 = models.resnet18()
- 2、修改全连接层的输出
num_ftrs = resnet18.fc.in_features
resnet18.fc = nn.Linear(num_ftrs, 4)#这里我们的模型分为 4 类
- 3、加载模型参数
checkpoint = torch.load(m_path)
resnet18.load_state_dict(checkpoint['model_state_dict'])
-
4、resnet18预训练模型下载
wget “https://download.pytorch.org/models/resnet18-5c106cde.pth”
或者：
浏览器下载：https://download.pytorch.org/models/resnet18-5c106cde.pth
使用预训练的 ResNet 18 快速训练自己的分类模型详细过程
- 1、依赖包
先导入依赖包：
import os
import torch.utils.data as data
import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim import lr_scheduler
from torchvision import datasets, models, transforms
from PIL import Image
import numpy as np
import cv2
import sys
- 2、加载预训练的模型
model_ft = models.resnet18(pretrained=True)
num_fits = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层 NUMCLASS=4
- 3、设置 device
device = torch.device('cuda:0') # 默认使用 GPU
model_ft = model_ft.to(device)
model_ft.cuda()
- 4、定义优化器
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
- 5、训练模型
num_epochs = 100
model_ft,arr_acc = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs)
- 6、保存模型
torch.save(model_ft.state_dict(), './model/my_model.pth')
- 7、定义自己的数据类
def default_loader(path):
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('RGB')
class CustomImageLoader(data.Dataset): # 定义自己的数据类
##自定义类型数据输入
def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader):
im_list = []
im_labels = []
with open(txt_path, 'r') as files:
for line in files:
#/x/y/a.jpg 1
#/x/y/b.jpg 2
items = line.split()
im_list.append(items[0])
im_labels.append(int(items[1]))
self.imgs = im_list
self.labels = im_labels
self.data_tranforms = data_transforms
self.loader = loader
self.dataset = dataset
def __len__(self):
return len(self.imgs)
def __getitem__(self, item):
img_name = self.imgs[item]
label = self.labels[item]
img = self.loader(img_name)
if self.data_tranforms is not None:
try:
img = self.data_tranforms[self.dataset](img)
except:
print("Cannot transform image: {}".format(img_name))
return img, label
数据转换器：根据自己的需要设置，比如下面的 64、50 是根据我的输入图像尺寸决定的。
resnet18中的默认值：
transforms.Resize(256),
transforms.CenterCrop(224),
我的data_tranforms：
data_tranforms={
'Train':transforms.Compose([
transforms.RandomResizedCrop(50), # 随机裁剪为不同的大小和宽高比,缩放所为制定的大小
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
]),
'Test':transforms.Compose([
transforms.Resize(64), # 变换大小
transforms.CenterCrop(50), # 中心裁剪
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
])
}
- 8、定义训练函数
image_datasets = {x : CustomImageLoader('/', # 默认目录为根目录，配搭文件中使用全路径
txt_path=('./data/{0}.txt'.format(x)), # 标签文件
data_transforms=data_tranforms,
dataset=x) for x in ['Train', 'Test']
}
dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True) for x in ['Train', 'Test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['Train', 'Test']} # 数据大小
def train_model(model, crtiation, optimizer,schedular, num_epochs=NUM_EPOCH):
begin_time = time.time()
best_weights = copy.deepcopy(model.state_dict())#copy the weights from the model
best_acc = 0.0
arr_acc = [] # 用于作图
for epoch in range(num_epochs):
print("-*-" * 20)
item_acc = []
for phase in ['Train', 'Test']:
if phase=='Train':
schedular.step()
model.train()
else:
model.eval()
running_loss = 0.0
running_acc = 0.0
for images, labels in dataloders[phase]:
images.to(device)
labels.to(device)
optimizer.zero_grad()
with torch.set_grad_enabled(phase=='Train'):
opt = model(images.cuda())
# opt = model(images)
_,pred = torch.max(opt,1)
labels = labels.cuda()
loss = crtiation(opt, labels)
if phase=='Train':
loss.backward()
optimizer.step()
running_loss += loss.item()*images.size(0)
running_acc += torch.sum(pred==labels)
epoch_loss = running_loss/dataset_sizes[phase]
epoch_acc = running_acc.double()/dataset_sizes[phase]
print('epoch={}, Phase={}, Loss={:.4f}, ACC:{:.4f}'.format(epoch, phase,
epoch_loss, epoch_acc))
item_acc.append(epoch_acc)
if phase == 'Test' and epoch_acc>best_acc:
# Upgrade the weights
best_acc=epoch_acc
best_weights = copy.deepcopy(model.state_dict())
arr_acc.append(item_acc)
time_elapes = time.time() - begin_time
print('Training Complete in{:.0f}m {:0f}s'.format(
time_elapes // 60, time_elapes % 60
))
print('Best Val ACC: {:}'.format(best_acc))
model.load_state_dict(best_weights) # 保存最好的参数
return model,arr_acc
完整代码
train.py：
import os
import torch.utils.data as data
import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim import lr_scheduler
from torchvision import datasets, models, transforms
from PIL import Image
import time
import copy
#import pandas as pd
#import matplotlib.pyplot as plt
import numpy as np
#%matplotlib inline
NUM_EPOCH = 100 # 默认迭代次数
batch_size = 64
device = torch.device('cuda:0') # 默认使用 GPU
NUMCLASS = 4 # 类别数
def default_loader(path):
with open(path, 'rb') as f:
with Image.open(f) as img:
return img.convert('RGB')
class CustomImageLoader(data.Dataset): # 定义自己的数据类
##自定义类型数据输入
def __init__(self, img_path, txt_path, dataset = '', data_transforms=None, loader = default_loader):
im_list = []
im_labels = []
with open(txt_path, 'r') as files:
for line in files:
#/x/y/a.jpg 1
#/x/y/b.jpg 2
items = line.split()
im_list.append(items[0])
im_labels.append(int(items[1]))
self.imgs = im_list
self.labels = im_labels
self.data_tranforms = data_transforms
self.loader = loader
self.dataset = dataset
def __len__(self):
return len(self.imgs)
def __getitem__(self, item):
img_name = self.imgs[item]
label = self.labels[item]
img = self.loader(img_name)
if self.data_tranforms is not None:
try:
img = self.data_tranforms[self.dataset](img)
except:
print("Cannot transform image: {}".format(img_name))
return img, label
data_tranforms={
'Train':transforms.Compose([
transforms.RandomResizedCrop(50), # 随机裁剪为不同的大小和宽高比,缩放所为制定的大小
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) # 各通道颜色的均值和方差,用于归一化
]),
'Test':transforms.Compose([
transforms.Resize(64), # 变换大小
transforms.CenterCrop(50), # 中心裁剪
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])
])
}
image_datasets = {x : CustomImageLoader('/', # 默认目录为根目录，配搭文件中使用全路径
txt_path=('./data/{0}.txt'.format(x)), # 标签文件
data_transforms=data_tranforms,
dataset=x) for x in ['Train', 'Test']
}
dataloders = {x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True) for x in ['Train', 'Test']}
dataset_sizes = {x: len(image_datasets[x]) for x in ['Train', 'Test']} # 数据大小
def train_model(model, crtiation, optimizer,schedular, num_epochs=NUM_EPOCH):
begin_time = time.time()
best_weights = copy.deepcopy(model.state_dict())#copy the weights from the model
best_acc = 0.0
arr_acc = [] # 用于作图
for epoch in range(num_epochs):
print("-*-" * 20)
item_acc = []
for phase in ['Train', 'Test']:
if phase=='Train':
schedular.step()
model.train()
else:
model.eval()
running_loss = 0.0
running_acc = 0.0
for images, labels in dataloders[phase]:
images.to(device)
labels.to(device)
optimizer.zero_grad()
with torch.set_grad_enabled(phase=='Train'):
opt = model(images.cuda())
# opt = model(images)
_,pred = torch.max(opt,1)
labels = labels.cuda()
loss = crtiation(opt, labels)
if phase=='Train':
loss.backward()
optimizer.step()
running_loss += loss.item()*images.size(0)
running_acc += torch.sum(pred==labels)
epoch_loss = running_loss/dataset_sizes[phase]
epoch_acc = running_acc.double()/dataset_sizes[phase]
print('epoch={}, Phase={}, Loss={:.4f}, ACC:{:.4f}'.format(epoch, phase,
epoch_loss, epoch_acc))
item_acc.append(epoch_acc)
if phase == 'Test' and epoch_acc>best_acc:
# Upgrade the weights
best_acc=epoch_acc
best_weights = copy.deepcopy(model.state_dict())
arr_acc.append(item_acc)
time_elapes = time.time() - begin_time
print('Training Complete in{:.0f}m {:0f}s'.format(
time_elapes // 60, time_elapes % 60
))
print('Best Val ACC: {:}'.format(best_acc))
model.load_state_dict(best_weights) # 保存最好的参数
return model,arr_acc
if __name__ == '__main__':
model_ft = models.resnet18(pretrained=True)
num_fits = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_fits, NUMCLASS) # 替换最后一个全连接层
model_ft = model_ft.to(device)
model_ft.cuda()
criterion = nn.CrossEntropyLoss()
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)
exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=10, gamma=0.1)
model_ft,arr_acc = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, NUM_EPOCH)
## 保存模型
torch.save(model_ft.state_dict(), './model/my_model.pth')