首页 » 机器视觉 » 正文
机器视觉

单发多框检测 SSD 代码测试

1. 说明

2. 训练SSD

2.1 训练SSD代码

# -*- coding: utf-8 -*-
# @Project: zc
# @Author: ZhangChu
# @File name: cls_predictor
# @Create time: 2022/1/2 16:37

# 1.导入相关数据库
import matplotlib.pyplot as plt
import torch
import torchvision
from torch import nn
from d2l import torch as d2l


# 2.类别预测卷积层
# 输入锚框数 -> 锚框数 * (类别数+1)
# num_inputs=8, num_anchors=5, num_classes=10
# 8 -> 5*11=55
def cls_predictor(num_inputs, num_anchors, num_classes):
	return nn.Conv2d(num_inputs, num_anchors * (num_classes + 1),
					 kernel_size=3, padding=1)


# 3.边缘框预测卷积层：每个锚框预测4个偏移量
# 输入锚框数 -> 锚框数*4
def bbox_predictor(num_inputs, num_anchors):
	return nn.Conv2d(num_inputs, num_anchors * 4,
					 kernel_size=3, padding=1)


def forward(x, block):
	return block(x)


# 5.Y1 = (2,55,20,20)
# 注：因为输入为(2,8,20,20)所以输入通道数为8
# cls_predictor的通道数是 num_input -> num_anchors*(num_classes+1)
# 8 -> 5*(10+1)=55 即输出通道数为55
# 故 Y1 = (2,55,20,20)
# Y2 = (2,33,10,10)
# 形状:[batch_size,num_channels,height,width]
Y1 = forward(torch.zeros((2, 8, 20, 20)), cls_predictor(8, 5, 10))
Y2 = forward(torch.zeros((2, 16, 10, 10)), cls_predictor(16, 3, 10))


# 6.将张量拉平，permute调换维度的顺序
# 输入形状:[batch_size,num_channels,height,width]
# 输出形状:[batch_size,height*width*num_channels]
def flatten_pred(pred):
	return torch.flatten(pred.permute(0, 2, 3, 1), start_dim=1)


# 7. 将张量拍平后叠加起来
def concat_preds(preds):
	return torch.cat([flatten_pred(p) for p in preds], dim=1)


# Y1 = [2,55,20,20]
# Y2 = [2,33,10,10]
# shape1 = [2,25300]
# 1.将Y1拉平得到[2,55*20*20]=[2,22000]
# 2.将Y2拉平得到[2,33*10*10]=[2,3300]
# 3.将Y1,Y2叠加得到shape1=[2,22000+3300]=[2,25300]
# shape1 = concat_preds([Y1, Y2])
#
# def down_sample_blk(in_channels, out_channels):
# 	blk = []
# 	for _ in range(2):
# 		blk.append(nn.Conv2d(in_channels, out_channels,
# 							 kernel_size=3, padding=1))
# 		blk.append(nn.BatchNorm2d(out_channels))
# 		blk.append(nn.ReLU())
# 		in_channels = out_channels
# 	blk.append(nn.MaxPool2d(2))
# 	return nn.Sequential(*blk)


# 8. 定义下采样的块block
# in_channels=3,out_channels=10
# blk =[
# 0:nn.Conv2d(3,10)
# 1:nn.BatchNorm2d(10)
# 2:nn.ReLU()
# 3:nn.Conv2d(3.10)
# 4:nn.BatchNorm2d()
# 5:nn.ReLU()
# 6:nn.MaxPool2d()]

def down_sample_blk(in_channels, out_channels):
	blk = []
	for _ in range(2):
		blk.append(nn.Conv2d(in_channels, out_channels,
							 kernel_size=3, padding=1))
		blk.append(nn.BatchNorm2d(out_channels))
		blk.append(nn.ReLU())
		in_channels = out_channels
	blk.append(nn.MaxPool2d(2))
	return nn.Sequential(*blk)


# m = down_sample_blk(3, 10)

# down_sample_blk 的作用是进行卷积计算，将输入通道数3，变成输出通道数10，
# 经过 nn.MaxPool2d 将尺寸从 (20 X 20) 变成 (10 X 10)
# 最后得到 x 的形状为 (2,10,10,10)
# x = forward(torch.zeros((2, 3, 20, 20)), down_sample_blk(3, 10)).shape
# x = [2,10,10,10]

# 9. 定义基础网络
def base_net():
	blk = []
	num_filters = [3, 16, 32, 64]
	for i in range(len(num_filters) - 1):
		blk.append(down_sample_blk(num_filters[i], num_filters[i + 1]))
	return nn.Sequential(*blk)


# 因为base_net的作用是进行了 3 次的down_sample_blk，每次就是将通道数进行改变，尺寸减半
# 故通道数变换如下 3   ->  16 -> 32 -> 64
#     尺寸变换如下 256 -> 128 -> 64 -> 32
# y 的形状的定义如下 :[batch_size,num_channels,height,width]
# 故 y 的形状变换如下: [2,3,256,256] -> [2,64,32,32]
# y = forward(torch.zeros((2, 3, 256, 256)), base_net()).shape
# print(f'y={y}')

# 10.获取网络块
def get_blk(i):
	if i == 0:
		blk = base_net()
	elif i == 1:
		blk = down_sample_blk(64, 128)
	elif i == 4:
		blk = nn.AdaptiveMaxPool2d((1, 1))
	else:
		blk = down_sample_blk(128, 128)
	return blk


# 11.定义block块前向传播函数
def blk_forward(X, blk, size, ratio, cls_predictor, bbox_predictor):
	"""

	:param X: 输入
	:param blk: 网络块
	:param size: size 缩放比
	:param ratio: 宽高比
	:param cls_predictor:类别分类器
	:param bbox_predictor: 边缘框分类器
	:return:(Y,anchors,cls_preds,bbox_preds)
	Y:输出；anchors：锚框，cls_preds:类别，bbox_preds：边缘框
	"""

	Y = blk(X)
	anchors = d2l.multibox_prior(Y, sizes=size, ratios=ratio)
	cls_preds = cls_predictor(Y)
	bbox_preds = bbox_predictor(Y)
	return (Y, anchors, cls_preds, bbox_preds)

# 12.初始化缩放比
sizes = [[0.2, 0.272], [0.37, 0.447], [0.54, 0.619], [0.71, 0.79],
		 [0.88, 0.961]]
ratios = [[1, 2, 0.5]] * 5
num_anchors = len(sizes[0]) + len(ratios[0]) - 1


# 13.定义小型的SSD神经网络
class TinySSD(nn.Module):
	def __init__(self, num_classes, **kwargs):
		super(TinySSD, self).__init__(**kwargs)
		self.num_classes = num_classes
		idx_to_in_channels = [64, 128, 128, 128, 128]
		for i in range(5):
			# 即赋值语句self.blk_i=get_blk(i)
			setattr(self, f'blk_{i}', get_blk(i))
			setattr(self, f'cls_{i}', cls_predictor(idx_to_in_channels[i],
													num_anchors, num_classes))
			setattr(self, f'bbox_{i}', bbox_predictor(idx_to_in_channels[i],
													  num_anchors))

	def forward(self, X):
		anchors, cls_preds, bbox_preds = [None] * 5, [None] * 5, [None] * 5
		for i in range(5):
			# getattr(self,'blk_%d'%i)即访问self.blk_i
			X, anchors[i], cls_preds[i], bbox_preds[i] = blk_forward(
				X, getattr(self, f'blk_{i}'), sizes[i], ratios[i],
				getattr(self, f'cls_{i}'), getattr(self, f'bbox_{i}'))
		anchors = torch.cat(anchors, dim=1)
		cls_preds = concat_preds(cls_preds)
		cls_preds = cls_preds.reshape(
			cls_preds.shape[0], -1, self.num_classes + 1)
		bbox_preds = concat_preds(bbox_preds)
		return anchors, cls_preds, bbox_preds


# 14.实例化TinySSD网络
net = TinySSD(num_classes=1)
X = torch.zeros((32, 3, 256, 256))
anchors, cls_preds, bbox_preds = net(X)

# print('output anchors:', anchors.shape)
# print('output class preds:', cls_preds.shape)
# print('output bbox preds:', bbox_preds.shape)

# 15.批量大小
batch_size = 32
# 16.训练集迭代器
train_iter, _ = d2l.load_data_bananas(batch_size)

# 17. device:GPU,net=TinySSD
device, net = d2l.try_gpu(), TinySSD(num_classes=1)

# 18. 定义优化器SGD
trainer = torch.optim.SGD(net.parameters(), lr=0.2, weight_decay=5e-4)

# 19.定义类别分类器的损失函数为交叉熵损失,reduction='none'：分别罗列出每个损失
cls_loss = nn.CrossEntropyLoss(reduction='none')

# 20.定义锚框分类器的损失函数为L1损失，reduction='none'：分别罗列出每个损失
bbox_loss = nn.L1Loss(reduction='none')


# 21.计算损失，类别损失cls,锚框损失bbox
def calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels, bbox_masks):
	batch_size, num_classes = cls_preds.shape[0], cls_preds.shape[2]
	cls = cls_loss(cls_preds.reshape(-1, num_classes),
				   cls_labels.reshape(-1)).reshape(batch_size, -1).mean(dim=1)
	bbox = bbox_loss(bbox_preds * bbox_masks,
					 bbox_labels * bbox_masks).mean(dim=1)
	return cls + bbox


# 22. 类别评估
def cls_eval(cls_preds, cls_labels):
	# 由于类别预测结果放在最后一维，argmax需要指定最后一维。
	return float((cls_preds.argmax(dim=-1).type(
		cls_labels.dtype) == cls_labels).sum())


# 23.锚框评估
def bbox_eval(bbox_preds, bbox_labels, bbox_masks):
	return float((torch.abs((bbox_labels - bbox_preds) * bbox_masks)).sum())


# 24. 定义训练次数 num_epochs,计时器 timer
num_epochs, timer = 20, d2l.Timer()

# 25. animator可视化
animator = d2l.Animator(xlabel='epoch', xlim=[1, num_epochs],
						legend=['class error', 'bbox mae'])
# 26. 将神经网络 net 放到 GPU上
net = net.to(device)

# 27. 开始迭代
for epoch in range(num_epochs):
	# 训练精确度的和，训练精确度的和中的示例数
	# 绝对误差的和，绝对误差的和中的示例数
	metric = d2l.Accumulator(4)
	# 开启网络训练模式
	net.train()
	for features, target in train_iter:
		# 将训练集的特征和标签提取出来
		timer.start()
		# 优化器梯度清零
		trainer.zero_grad()
		# 将训练集中的 features 给 X,标签target 给 Y ，数据一起放在 GPU上
		X, Y = features.to(device), target.to(device)
		# 生成多尺度的锚框，为每个锚框预测类别和偏移量
		# X 在给定的网络中得到了锚框，类别预测，锚框偏移量预测
		anchors, cls_preds, bbox_preds = net(X)
		# 为每个锚框标注类别和偏移量
		bbox_labels, bbox_masks, cls_labels = d2l.multibox_target(anchors, Y)
		# 根据类别和偏移量的预测和标注值计算损失函数
		l = calc_loss(cls_preds, cls_labels, bbox_preds, bbox_labels,
					  bbox_masks)
		# 将损失函数求均值，并反向传播
		l.mean().backward()
		# 更新优化器
		trainer.step()
		# 将评估值画出来
		metric.add(cls_eval(cls_preds, cls_labels), cls_labels.numel(),
				   bbox_eval(bbox_preds, bbox_labels, bbox_masks),
				   bbox_labels.numel())
	cls_err, bbox_mae = 1 - metric[0] / metric[1], metric[2] / metric[3]
	animator.add(epoch + 1, (cls_err, bbox_mae))
print(f'class err {cls_err:.2e}, bbox mae {bbox_mae:.2e}')
print(f'{len(train_iter.dataset) / timer.stop():.1f} examples/sec on '
	  f'{str(device)}')

plt.show()

2.2 训练SSD结果

output anchors: torch.Size([1, 5444, 4])
output class preds: torch.Size([32, 5444, 2])
output bbox preds: torch.Size([32, 21776])

class err 3.38e-03, bbox mae 3.26e-03
4123.0 examples/sec on cuda:0

在这里插入图片描述

3. 预测SSD

3.1 预测 SSD 代码

# 28.读取需要预测的图片
X = torchvision.io.read_image('D:/zc/img/banana.jpg').unsqueeze(0).float()
img = X.squeeze(0).permute(1, 2, 0).long()


# 29.预测
def predict(X):
	# 网络预测模式
	net.eval()
	# X通过网络得到预测值
	anchors, cls_preds, bbox_preds = net(X.to(device))
	# 求出类别最大的概率
	cls_probs = F.softmax(cls_preds, dim=2).permute(0, 2, 1)
	# 为图片输出多锚框
	output = d2l.multibox_detection(cls_probs, bbox_preds, anchors)
	idx = [i for i, row in enumerate(output[0]) if row[0] != -1]
	return output[0, idx]


output = predict(X)


# 30.显示图片
def display(img, output, threshold):
	# 设置画布大小
	d2l.set_figsize((5, 5))
	# 显示图片
	fig = d2l.plt.imshow(img)
	# 将大于指定置信度的锚框输出
	for row in output:
		score = float(row[1])
		if score < threshold:
			continue
		h, w = img.shape[0:2]
		bbox = [row[2:6] * torch.tensor((w, h, w, h), device=row.device)]
		d2l.show_bboxes(fig.axes, bbox, '%.2f' % score, 'w')

# 31.显示图片
display(img, output.cpu(), threshold=0.9)

plt.show()

3.2 预测 SSD 结果

在这里插入图片描述

版权声明：本文为CSDN博主「取个名字真难呐」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/scar2016/article/details/122280762

标签： · · · ·
点赞 0

我还没有学会写个人说明！

暂无评论

发表评论

相关推荐

yolov5训练数据集划分
机器视觉 2024-11-22

yolov5训练数据集划分

yolov5训练数据集划分 按照默认8：1：1划分训练集，测试集，验证集。 txt文件出现在imageset文件夹。 import os import randomtrainval_pe

分享到：