轻松学Pytorch-迁移学习实现表面缺陷检查
引言 ·
大家好,我又又好久没有更新这个系列了,主要原因归根结底只有一个懒,所谓 一勤天下无难事,百思心中有良谋。以后还争取每周更新,这次隔了一周没有更新,对不起大家了。今天给大家更新的是如何基于torchvision自带的模型完成图像分类任务的迁移学习,前面我们已经完成了对对象检测任务的迁移学习,这里补上针对图像分类任务的迁移学习,官方的文档比较啰嗦,看了之后其实可操作性很低,特别是对于初学者,估计看了之后就发懵的那种。本人重新改写了一波,代码简洁易懂,然后把训练结果导出ONNX,使用OpenCV DNN调用部署,非常实用!废话不多说了,少吹水!
数据集
东北大学热轧带钢表面缺陷数据集,该数据集是东北大学的宋克臣等几位老师收集的,一共包含了三类数据。这里使用(NEU surface defect database),数据集收集了夹杂、划痕、压入氧化皮、裂纹、麻点和斑块总计6种缺陷,每种缺陷300张,图像尺寸为200×200。部分示例如下:
基于该数据集,实现pytorch数据类,完成数据集的加载与预处理的代码如下:
class SurfaceDefectDataset(Dataset):
def __init__(self, root_dir):
self.transform = transforms.Compose([transforms.ToTensor()])
img_files = os.listdir(root_dir)
self.defect_types = []
self.images = []
index = 0
for file_name in img_files:
defect_attrs = file_name.split("_")
d_index = defect_labels.index(defect_attrs[0])
self.images.append(os.path.join(root_dir, file_name))
self.defect_types.append(d_index)
index += 1
def __len__(self):
return len(self.images)
def num_of_samples(self):
return len(self.images)
def __getitem__(self, idx):
if torch.is_tensor(idx):
idx = idx.tolist()
image_path = self.images[idx]
else:
image_path = self.images[idx]
img = cv.imread(image_path) # BGR order
h, w, c = img.shape
# rescale
img = cv.resize(img, (200, 200))
img = (np.float32(img) /255.0 - 0.5) / 0.5
# H, W C to C, H, W
img = img.transpose((2, 0, 1))
sample = {'image': torch.from_numpy(img), 'defect': self.defect_types[idx]}
return sample
怎么下载该数据集,后台回复"NEU"关键字即可获取下载地址
模型使用
Pytorchvison支持多种图像分类模型,这里我们选择残差网络模型作为迁移学习的基础模型,对输出层(最后一层)改为六个类别,其它特征层选择在训练时候微调参数。常见的ResNet网络模型如下:
基于ResNet18完成网络模型修改,最终的模型实现代码如下:
class SurfaceDefectResNet(torch.nn.Module):
def __init__(self):
super(SurfaceDefectResNet, self).__init__()
self.cnn_layers = torchvision.models.resnet18(pretrained=True)
num_ftrs = self.cnn_layers.fc.in_features
self.cnn_layers.fc = torch.nn.Linear(num_ftrs, 6)
def forward(self, x):
# stack convolution layers
out = self.cnn_layers(x)
return out
模型训练与测试
模型训练跟前面讲的一些图像分类模型训练方式并无不同,基于交叉熵损失,完成训练,每个批次4张图像或者8张图,训练15个epoch之后,保存模型。然后使用模型测试35张测试图像,发现有两张预测错误,其余均正确。训练模型的代码如下:
# 训练模型的次数
num_epochs = 15
# optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
model.train()
# 损失函数
cross_loss = torch.nn.CrossEntropyLoss()
index = 0
for epoch in range(num_epochs):
train_loss = 0.0
for i_batch, sample_batched in enumerate(dataloader):
images_batch, label_batch =
sample_batched['image'], sample_batched['defect']
if train_on_gpu:
images_batch, label_batch= images_batch.cuda(), label_batch.cuda()
optimizer.zero_grad()
# forward pass: compute predicted outputs by passing inputs to the model
m_label_out_ = model(images_batch)
label_batch = label_batch.long()
# calculate the batch loss
loss = cross_loss(m_label_out_, label_batch)
# backward pass: compute gradient of the loss with respect to model parameters
loss.backward()
# perform a single optimization step (parameter update)
optimizer.step()
# update training loss
train_loss += loss.item()
if index % 100 == 0:
print('step: {} tTraining Loss: {:.6f} '.format(index, loss.item()))
index += 1
# 计算平均损失
train_loss = train_loss / num_train_samples
# 显示训练集与验证集的损失函数
print('Epoch: {} tTraining Loss: {:.6f} '.format(epoch, train_loss))
# save model
model.eval()
torch.save(model, 'surface_defect_model.pt')
转为为ONNX模式,OpenCV DNN部署调用,代码如下:
defect_net = cv.dnn.readNetFromONNX("surface_defect_resnet18.onnx")
root_dir = "D:/pytorch/enu_surface_defect/test"
fileNames = os.listdir(root_dir)
for f in fileNames:
image = cv.imread(os.path.join(root_dir, f))
blob = cv.dnn.blobFromImage(image, 0.00392, (200, 200), (127, 127, 127)) / 0.5
defect_net.setInput(blob)
res = defect_net.forward()
idx = np.argmax(np.reshape(res, (6)))
defect_txt = defect_labels[idx]
cv.putText(image, defect_txt, (10, 25), cv.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
cv.imshow("input", image)
print(f, defect_txt)
cv.waitKey(0)
cv.destroyAllWindows()
预测运行结果如下:
运行结果与pytorch调用模型运行结果保持一致。由于这个是一个专栏,很多代码在以前的文章中已经给出了,这里就没有重复贴代码!
- ASP.NET Identity 2新增双重认证、帐号锁定、防伪印章功能并修复了一些bug
- Serverless 应用开发指南:基于 Serverless 与 Lambda 的微信公共平台
- Serverless 应用开发指南:serverless 的 hello, world
- 通过一组RESTful API暴露CQRS系统功能
- 通过使用结构化数据 JSON-LD,我为网站带来了更多的流量
- 使用 OWIN Self-Host ASP.NET Web API 2
- c#开源消息队列中间件EQueue 教程
- Serverless 框架 OpenWhisk 开发指南:使用 Node.js 编写 hello, world
- GOTO Berlin: Web API设计原则
- 使用 ServiceStack 构建跨平台 Web 服务
- 使用 OpenWhisk 自建 Serverless 服务
- 如何在 8 小时内开发上线一个在线表单系统
- 让Response.Redirect页面重定向更有效率
- 使用 adr 轻松创建 “程序员友好” 的轻量级文档
- JavaScript 教程
- JavaScript 编辑工具
- JavaScript 与HTML
- JavaScript 与Java
- JavaScript 数据结构
- JavaScript 基本数据类型
- JavaScript 特殊数据类型
- JavaScript 运算符
- JavaScript typeof 运算符
- JavaScript 表达式
- JavaScript 类型转换
- JavaScript 基本语法
- JavaScript 注释
- Javascript 基本处理流程
- Javascript 选择结构
- Javascript if 语句
- Javascript if 语句的嵌套
- Javascript switch 语句
- Javascript 循环结构
- Javascript 循环结构实例
- Javascript 跳转语句
- Javascript 控制语句总结
- Javascript 函数介绍
- Javascript 函数的定义
- Javascript 函数调用
- Javascript 几种特殊的函数
- JavaScript 内置函数简介
- Javascript eval() 函数
- Javascript isFinite() 函数
- Javascript isNaN() 函数
- parseInt() 与 parseFloat()
- escape() 与 unescape()
- Javascript 字符串介绍
- Javascript length属性
- javascript 字符串函数
- Javascript 日期对象简介
- Javascript 日期对象用途
- Date 对象属性和方法
- Javascript 数组是什么
- Javascript 创建数组
- Javascript 数组赋值与取值
- Javascript 数组属性和方法