classDetectionModel(nn.Module):
def__init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classessuper().__init__()
...
# 假如直接传入的是dict则无需处理; 假如不是则运用yaml.safe_load加载yaml文件withopen(cfg, errors='ignore') as f:
self.yaml = yaml.safe_load(f) # model dict
...
# 创立网络模型# self.model: 初始化的整个网络模型(包括Detect层结构)# self.save: 一切层结构中from不等于-1的序号,并排好序 [4, 6, 10, 14, 17, 20, 23]
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
...
defforward(self, x, augment=False, profile=False, visualize=False): # debug同样需求第三次才干正常跳进来if augment: # use Test Time Augmentation(TTA), 假如打开会对图片进行scale和flipreturn self._forward_augment(x) # augmented inference, Nonereturn self._forward_once(x, profile, visualize) # single-scale inference, train# 运用TTA进行推理(当然仍是会调用普通推理完成前向传达)def_forward_augment(self, x):
img_size = x.shape[-2:] # height, width
s = [1, 0.83, 0.67] # scales
f = [None, 3, None] # flips (2-ud上下flip, 3-lr左右flip)
y = [] # outputs# 这儿相当于对输入x进行3次不同参数的测验数据增强推理, 每次的推理结构都保存在列表y中for si, fi inzip(s, f):
# scale_img缩放图片尺度# 经过普通的双线性插值完成,依据ratio来操控图片的缩放比例,最终经过pad 0补齐到原图的尺度
xi = scale_img(x.flip(fi) if fi else x, si, gs=int(self.stride.max()))
yi = self._forward_once(xi)[0] # forward:torch.Size([1, 25200, 25])# cv2.imwrite(f'img_{si}.jpg', 255 * xi[0].cpu().numpy().transpose((1, 2, 0))[:, :, ::-1]) # save# _descale_pred将推理成果康复到相对原图图片尺度, 只对坐标xywh:yi[..., :4]进行康复# 假如f=2,进行上下翻转; 假如f=3,进行左右翻转
yi = self._descale_pred(yi, fi, si, img_size)
y.append(yi) # [b, 25200, 25] / [b, 18207, 25] / [b, 12348, 25]# 把第一层的后面一部分的猜测成果去掉, 也把最终一层的前面一部分的猜测成果去掉# [b, 24000, 25] / [b, 18207, 25] / [b, 2940, 25]# 筛除的可能是重复的部分吧, 进步运转速度(有了解的朋友请告诉我一下)
y = self._clip_augmented(y) # clip augmented tailsreturn torch.cat(y, 1), None# augmented inference, train# 普通推理def_forward_once(self, x, profile=False, visualize=False):
# y列表用来保存中间特征图; dt用来记载每个模块履行10次的均匀时长
y, dt = [], [] # outputs# 对sequence模型进行遍历操作, 不断地对输入x进行处理, 中间成果需求保存的时分别的存储到列表y中for m in self.model:
# 假如仅仅对前一个模块的输出进行操作, 则需求提取直接保存的中间特征图进行操作,# 一般是concat处理, 对当前层与之前曾进行一个concat再卷积; detect模块也需求提取3个特征层来处理if m.f != -1: # if not from previous layer
x = y[m.f] ifisinstance(m.f, int) else [x if j == -1else y[j] for j in m.f] # from earlier layers# profile参数打开会记载每个模块的均匀履行10次的时长和flops用于剖析模型的瓶颈, 进步模型的履行速度和下降显存占用if profile:
self._profile_one_layer(m, x, dt)
# 运用当前模块对特征图进行处理# 假如是concat模块: 则x是一个特征图列表, 则对其进行拼接处理, 再交给下一个卷积模块;# 假如是C3, Conv等普通的模块: 则x是单一特征图# 假如是detct模块: 则x是3个特征图的列表 (练习与推理回来的内容不一样)
x = m(x) # run# self.save: 把一切层结构中from不是-1的值记下并排序 [4, 6, 10, 14, 17, 20, 23]
y.append(x if m.i in self.save elseNone) # save output# 特征可视化if visualize:
feature_visualization(x, m.type, m.i, save_dir=visualize)
return x
# 翻转数据增强def_descale_pred(self, p, flips, scale, img_size):
# de-scale predictions following augmented inference (inverse operation)if self.inplace:
p[..., :4] /= scale # de-scale xywh坐标缩放回原来巨细# f=2,进行上下翻转if flips == 2:
p[..., 1] = img_size[0] - p[..., 1] # de-flip ud# f=3,进行左右翻转elif flips == 3:
p[..., 0] = img_size[1] - p[..., 0] # de-flip lrelse:
x, y, wh = p[..., 0:1] / scale, p[..., 1:2] / scale, p[..., 2:4] / scale # de-scaleif flips == 2:
y = img_size[0] - y # de-flip udelif flips == 3:
x = img_size[1] - x # de-flip lr
p = torch.cat((x, y, wh, p[..., 4:]), -1)
return p
# 这儿y的一个包括3个子列表的列表, 经过对输入图画x进行了3次不同尺度的改换, 所以得到了3个inference结构# 这儿看不太懂, 不过大概做的工作便是对第一个列表与最终一个列表的成果做一些过滤处理# 把第一层的后面一部分的猜测成果去掉, 也把最终一层的前面一部分的猜测成果去掉, 然后剩余的concat为一个部分def_clip_augmented(self, y):
# Clip YOLOv5 augmented inference tails
nl = self.model[-1].nl # Detect(): number of detection layers (P3-P5)
g = sum(4 ** x for x inrange(nl)) # grid points
e = 1# exclude layer count
i = (y[0].shape[1] // g) * sum(4 ** x for x inrange(e)) # indices: (25200 // 21) * 1 = 1200
y[0] = y[0][:, :-i] # large: (1,25200,25) -> (1,24000,25)
i = (y[-1].shape[1] // g) * sum(4 ** (nl - 1 - x) for x inrange(e)) # indices: (12348 // 21) * 16 = 9408
y[-1] = y[-1][:, i:] # small: (1,12348,25) -> (1,2940,25)return y
...
# 经过普通的双线性插值完成,依据ratio来操控图片的缩放比例,最终经过pad 0补齐到原图的尺度defscale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416)# scales img(bs,3,y,x) by ratio constrained to gs-multipleif ratio == 1.0:
return img
else:
h, w = img.shape[2:]
s = (int(h * ratio), int(w * ratio)) # new size
img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resizeifnot same_shape: # pad/crop img
h, w = [math.ceil(x * ratio / gs) * gs for x in (h, w)]
return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean