yolo源码注释4——yolo-py
代码基于yolov5 v6.0
目录：
- yolo源码注释1——文件结构
- yolo源码注释2——数据集配置文件
- yolo源码注释3——模型配置文件
- yolo源码注释4——yolo-py
yolo.py 用于搭建 yolov5 的网络模型，主要包含 3 部分：
- Detect：Detect 层
- Model：搭建网络
- parse_model：根据配置实例化模块
Model（仅注释了 init 函数）：
class Model(nn.Module):
# YOLOv5 model
def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes
super().__init__()
if isinstance(cfg, dict):
self.yaml = cfg # model dict
else: # is *.yaml
import yaml
self.yaml_file = Path(cfg).name
with open(cfg, encoding='ascii', errors='ignore') as f:
self.yaml = yaml.safe_load(f)
# Define model
ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels
if nc and nc != self.yaml['nc']:
LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}")
self.yaml['nc'] = nc # override yaml value
if anchors:
LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}')
self.yaml['anchors'] = round(anchors) # override yaml value
# 根据配置搭建网络
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch])
self.names = [str(i) for i in range(self.yaml['nc'])] # default names
self.inplace = self.yaml.get('inplace', True)
# 计算生成 anchors 时的步长
m = self.model[-1] # Detect()
if isinstance(m, Detect):
s = 256 # 2x min stride
m.inplace = self.inplace
m.stride = torch.tensor([s / x.shape[-2] for x in self.forward(torch.zeros(1, ch, s, s))]) # forward
check_anchor_order(m) # must be in pixel-space (not grid-space)
m.anchors /= m.stride.view(-1, 1, 1)
self.stride = m.stride
self._initialize_biases() # only run once
# Init weights, biases
initialize_weights(self)
self.info()
LOGGER.info('')
parse_model：
def parse_model(d, ch): # model_dict, input_channels(3)
LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}")
anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple']
na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors
no = na * (nc + 5) # number of outputs = anchors * (classes + 5)
# layers： 保存每一层的结构
# save： 记录 from 不是 -1 的层，即需要多个输入的层如 Concat 和 Detect 层
# c2: 当前层输出的特征图数量
layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out
for i, (f, n, m, args) in enumerate(
d['backbone'] + d['head']): # from:-1, number:1, module:'Conv', args:[64, 6, 2, 2]
m = eval(m) if isinstance(m, str) else m # eval strings, m: