yolov5输出张量

### YOLOv5 输出张量格式解释

YOLOv5 的输出由多个不同尺度的特征图组成，这些特征图用于多尺度检测。对于每个尺度上的每一个网格单元格 (grid cell)，模型会预测三个边界框及其对应的类别概率和置信度。

#### 特征图结构
假设输入图像大小为 \(640 \times 640\) 像素，则对应三种不同的下采样比例下的特征图尺寸分别为：

- **大尺度**：\(80 \times 80\)
- **中尺度**：\(40 \times 40\)
- **小尺度**：\(20 \times 20\)

每种尺度上都有固定数量的锚点(anchor boxes), 默认情况下是3个[^1]。

#### 单个输出张量的具体构成
对于任意一层（比如上述提到的大、中、小任一层次），其输出张量形状可以表示为:

\[ B \times H \times W \times A \times (C + 5) \]

其中，
- \(B\) 表示批次(batch size);
- \(H, W\) 分别代表该层的高度(height) 和宽度(width)[^4];
- \(A=3\) 是指每个位置预设了3个先验框;
- \(C\) 则是指类别的数目；
- 数字`5` 对应于每个边框额外附带的信息——即中心坐标 (\(x,y\))、宽高(\(w,h\))以及存在目标的概率(objectness score).

因此，在实际应用中如果设定分类数为80的话，那么最终得到的就是形如 `[batch_size, height, width, num_anchors_per_scale, 85]` 这样的张量.

import torch

def get_yolov5_output_shape(input_height, input_width, batch_size, classes_count):
    scales = [(input_height//stride, input_width//stride) for stride in [8, 16, 32]]
    
    outputs_shapes = []
    anchors_per_scale = 3
    
    for h,w in scales:
        shape = (
            batch_size,
            h,
            w,
            anchors_per_scale,
            5 + classes_count # xywh + obj_confidence + class_probs
        )
        
        outputs_shapes.append(shape)
        
    return outputs_shapes


output_shapes = get_yolov5_output_shape(640, 640, 1, 80)
for i,shape in enumerate(output_shapes):
    print(f'Scale {i}:', ' '.join(map(str, shape)))