Posted on

2021-09-25 In machine learning

Disqus: 0 Comments

# 对于里面的元素
# 如果是元组，代表：(输出通道, 卷积核尺寸, 步长)
# YOLOv3中所有的卷积块（注意是卷积块，它由卷积层+批标准化层+LeakyReLU层构成）都是相同的，在下面的代码中用CNNBlock类实现
# 如果是列表，"B"代表残差块Residual Block，后面的次数代表重复次数，在下面用ResidualBlock类实现
# 如果是字符，那么"S"代表Scale不同尺度预测块，在此处计算损失，在下面用ScalePrediction类实现
# "U"代表Upsampling上采样，且与上一层进行连接，生成新的尺度预测
config = [
    (32, 3, 1),
    (64, 3, 2),
    ["B", 1],
    (128, 3, 2),
    ["B", 2],
    (256, 3, 2),
    ["B", 8],
    (512, 3, 2),
    ["B", 8],
    (1024, 3, 2),
    ["B", 4],  # 到这里就是Darknet-53 backbone，53是全部卷积层的个数，它会在imagenet上进行预训练
    (512, 1, 1),
    (1024, 3, 1),
    "S",
    (256, 1, 1),
    "U",
    (256, 1, 1),
    (512, 3, 1),
    "S",
    (128, 1, 1),
    "U",
    (128, 1, 1),
    (256, 3, 1),
    "S",
]

class CNNBlock(nn.Module):
    # 此处会加上一个BN层的开关，如果关了BN层，就相当于是只有卷积层，而不是卷积块
   # 不加BN和ReLU层的纯卷积层是会在尺度预测的地方用到，即网络末端的卷积是纯卷积层
   # 同时使用kwargs参数接收其他参数，比如kerneal size，stride，padding等参数
   # 在整个网络中图像的宽高变化即维度压缩，是通过卷积块的stride参数来实现的
   # 由下面的分析可知，在残差块中图像宽高不变，但两个残差块中间的卷积块的stride为2，此时会对图像的宽高进行压缩减半
   # 整个网络中压缩最厉害的分支是一共压缩了5次，即压缩了32倍，另外两支分别压缩了16倍和8倍
    def __init__(self, in_channels, out_channels, bn_act=True, **kwargs):
        super().__init__()
       # 如果使用了BN，那么偏置这个参数就没必要了，所以此处会根据BN层的有无进行偏置bias参数的开关
        self.conv = nn.Conv2d(in_channels, out_channels, bias=not bn_act, **kwargs)
        self.bn = nn.BatchNorm2d(out_channels)
        self.leaky = nn.LeakyReLU(0.1)
        self.use_bn_act = bn_act

def forward(self, x):
        if self.use_bn_act:
            return self.leaky(self.bn(self.conv(x)))
        else:
            return self.conv(x)

class ResidualBlock(nn.Module):
    # 这里给出了是否使用残差连接的开关，在darknet-53部分都是打开残差连接，但到了尺度预测部分，该残差块是关闭了残差连接
    def __init__(self, channels, use_residual=True, num_repeats=1):
        super().__init__()
        # 整个残差块是一个ModuleList
        self.layers = nn.ModuleList()
       # 根据重复次数进行循环
        for repeat in range(num_repeats):
            self.layers += [
                # 每个残差块中的第一个卷积块都是通道数减半，卷积核尺寸为1，步长是默认的1，填充是默认的0，因此图像的输入和输出宽高不变
            # 第二个卷积块通道数变为两倍，卷积核尺寸为3，填充为1，步长仍是默认的1，因此图像的输入和输出宽高不变
            # 所以，总的来说，经过一个残差块后，图像的通道数、宽和高都不会变
                nn.Sequential(
                    CNNBlock(channels, channels // 2, kernel_size=1),
                    CNNBlock(channels // 2, channels, kernel_size=3, padding=1),
                )
            ]

self.use_residual = use_residual
        self.num_repeats = num_repeats

def forward(self, x):
        for layer in self.layers:
            # 如果启用残差连接，那么就直接将x和经过处理后的x相加
            if self.use_residual:
                x = x + layer(x)
            # 如果没有启用残差连接，那么就直接处理x，不管作为输入的x
            else:
                x = layer(x)

return x

class ScalePrediction(nn.Module):
    def __init__(self, in_channels, num_classes):
        super().__init__()
        self.pred = nn.Sequential(
            # 在每个尺度预测块中，先用一个卷积块将通道数加倍，同时通过设置卷积核尺寸为3，填充为1，步长是默认的1，来保持宽高不变
            CNNBlock(in_channels, 2 * in_channels, kernel_size=3, padding=1),
            # 然后将得到的特征图通过一个卷积块转化为最终想要的向量的模样
         # 3指的是对于对于每一个grid cell，都有3个anchor boxes
         # 对于每一个anchor box，都需要有num_classes+5个元素，前面是类别数目，比如20，5是包含了x, y, w, h和置信度
         # 注意这里不使用BN层，同时卷积核为1，步长是默认的1，填充是默认的0，因此宽高不变
            CNNBlock(
                2 * in_channels, (num_classes + 5) * 3, bn_act=False, kernel_size=1
            ),
        )
        self.num_classes = num_classes

def forward(self, x):
        return (
            self.pred(x)
            # 对x进行预测后，需要对结果进行reshape，形状依次为batch size、3、类别数+5、特征图宽度、特征图高度
            .reshape(x.shape[0], 3, self.num_classes + 5, x.shape[2], x.shape[3])
            # 再交换一下维度，把宽、高提到(类别数+5)的前面
         # 比如某一个尺度预测后，得到的向量形状为N x 3 x 13 x 13 x (5+num_classes)，grid cell就是13x13大小
            .permute(0, 1, 3, 4, 2)
        )

class YOLOv3(nn.Module):
    # 输入通道默认为3， 类别数默认为80
    def __init__(self, in_channels=3, num_classes=80):
        super().__init__()
        self.num_classes = num_classes
        self.in_channels = in_channels
        self.layers = self._create_conv_layers()

def _create_conv_layers(self):
        # 将所有模型组件都放在ModuleList中
        layers = nn.ModuleList()
        in_channels = self.in_channels

# 开始解析上面的config配置
        for module in config:
            # 如果元素是个元组，代表它是个卷积块
            if isinstance(module, tuple):
                # 取出卷积块的相应配置
                out_channels, kernel_size, stride = module
                # 往整个网络里添加卷积块
                layers.append(
                    CNNBlock(
                        in_channels,
                        out_channels,
                        kernel_size=kernel_size,
                        stride=stride,
                        # 如果卷积核为3，则填充为1，否则就填充为0，这样是为了当卷积核为3、步长为2时，填充设为1，此时宽高减半
                  # Pytorch默认卷积层的尺寸计算是向下取整，即(k+2*1-3)/2+1=k/2+floor(-0.5)+1=k/2-1+1=k/2
                        padding=1 if kernel_size == 3 else 0,
                    )
                )
                # 更新通道数
                in_channels = out_channels

# 如果元素是个列表，代表是残差块
            elif isinstance(module, list):
                # 取出残差块的相应配置
                num_repeats = module[1]
                # 往整个网络里添加残差块
                layers.append(ResidualBlock(in_channels, num_repeats=num_repeats,))

# 如果元素是个字符，那么就进入尺度预测模块
            elif isinstance(module, str):
                # 如果是S，代表要进行在某一尺度上的预测了
                if module == "S":
                    layers += [
                        # 下面这三块的网络架构参考下面那张YOLOv3的架构图
                  # 原码中残差块只重复了1次，为了与下面架构图中的YoloBlock相对应，这里改为重复2次，影响不大，因为在残差块中不改变图像大小
                  # 同时注意此时残差块关闭了残差连接
                        ResidualBlock(in_channels, use_residual=False, num_repeats=2),
                        CNNBlock(in_channels, in_channels // 2, kernel_size=1),
                        ScalePrediction(in_channels // 2, num_classes=self.num_classes),
                    ]
                    # 更新一下通道数
                    in_channels = in_channels // 2

# 如果是U，则进入上采样
                elif module == "U":
                    layers.append(nn.Upsample(scale_factor=2),)
                    # 通道数变为3倍，原因是这个地方进行了通道连接concatenation操作
               # 特别注意的是，不要在这个地方推导图像在整个模型中的处理过程，因为此时会发现前后通道数是不符的，因为通道一下从256跳到了768
               # 这个地方不是forward函数，并不是真正的数据处理过程，可以理解成这个地方仅是模型架构定义
                    in_channels = in_channels * 3

return layers

def forward(self, x):
        # 每一个尺度下都有一个output，这里用一个列表来承载三个output
        outputs = []
        # 存放进入不同预测分支的中间计算结果
        route_connections = []
        # 对网络中的每一层进行遍历
        for layer in self.layers:
            # 如果是尺度预测层，表示进入某一尺度的预测阶段，即进入某一个预测分支
            if isinstance(layer, ScalePrediction):
                # 将预测结果添加进outputs中，注意这个地方是对x的一个分叉计算
            # 即x在这里走了两条路，一条路是进入尺度预测模块进行计算，另一条路是继续呆在主分支中，用于后续计算
                outputs.append(layer(x))
                # 返回到主分支中
                continue

# 对常规的网络层进行计算，包含卷积块和重复次数不为8的残差块
            x = layer(x)

# 对于重复次数为8的残差块，由架构图可知，都是在这里进入不同的尺度预测分支
            if isinstance(layer, ResidualBlock) and layer.num_repeats == 8:
                # 将需要进入某分支的结果存放起来
                route_connections.append(x)

# 如果是遇到上采样模块
            elif isinstance(layer, nn.Upsample):
                # 就会将当前x与存放中间结果的route中的最后一个中间结果进行连接concatenation
            # 这个地方会将通道数变为3倍，因为上采样后的图像为n_channel，原主分支中的图像为2*n_channel，连接后就变为3*n_channel
                x = torch.cat([x, route_connections[-1]], dim=1)
                # 用完最后一个元素就把它丢了，这样就能在下一次取到上一个存储的中间结果
                route_connections.pop()

# 最终outputs里是存放了三个尺度的预测模型
        return outputs

class YOLODataset(Dataset):
    def __init__(
        self,
        csv_file, # csv文件路径
        img_dir, # 图像文件路径
        label_dir, # 标签文件路径
        anchors, # 九个锚框
        image_size=416, # 图像尺寸
        S=[13, 26, 52], # 三个特征图大小
        C=20, # 类别数
        transform=None, # 图像变换
    ):
        self.annotations = pd.read_csv(csv_file) # 图像和标签成对出现
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.image_size = image_size
        self.transform = transform
        self.S = S
        self.anchors = torch.tensor(anchors[0] + anchors[1] + anchors[2])  # 将三个尺度下的三个锚框连起来，注意两个list相加就是join的效果
        self.num_anchors = self.anchors.shape[0]
        self.num_anchors_per_scale = self.num_anchors // 3
        self.C = C
        self.ignore_iou_thresh = 0.5 # 这个阈值用来区分忽略样例和负例，详见上面的解析

def __len__(self):
        return len(self.annotations)

def __getitem__(self, index):
        label_path = os.path.join(self.label_dir, self.annotations.iloc[index, 1]) # 1就是代表第二列，即标签列
        bboxes = np.roll(np.loadtxt(fname=label_path, delimiter=" ", ndmin=2), 4, axis=1).tolist() # 得到的边界框格式为(x, y, w, h, 类别)

# 读入图像文件
        img_path = os.path.join(self.img_dir, self.annotations.iloc[index, 0])
        image = np.array(Image.open(img_path).convert("RGB"))

# 图像增强变换
        if self.transform:
            augmentations = self.transform(image=image, bboxes=bboxes)
            image = augmentations["image"]
            bboxes = augmentations["bboxes"]

# 最终目标是在三个特征图尺度上，每个尺度的每个格点上都有(self.num_anchors // 3)个预测框，每个预测框上都有6个分量，即[置信度标签, x, y, w, h, 类别]
        # 置信度标签为1，代表正例；标签为-1，代表忽略样例；标签为0，代表负例。
        targets = [torch.zeros((self.num_anchors // 3, S, S, 6)) for S in self.S]

# 对该图像中的所有的边界框进行循环，目的是为了确定哪个锚框、哪个格点与其对应
        for box in bboxes: 
            # 确定与边界框对应的锚框是通过计算两者之间的IoU
            # box的第2、3元素就是宽度和高度
            # 这里的IoU计算相当于将锚框和边界框的中心放在一块，然后根据它们的宽高来计算
            # 即为了确定哪一种形状的锚框与该边界框最相近
            iou_anchors = iou(torch.tensor(box[2:4]), self.anchors)
            # 找出重合度最大的即最好的锚框
            anchor_indices = iou_anchors.argsort(descending=True, dim=0)
            # 取出该边界框的x, y, w, h
            x, y, width, height, class_label = box
            # 下面这个列表初始化为False，但最终目的是变为True，即保证在每个尺度下该边界框都有对应的锚框
            has_anchor = [False] * 3  # each scale should have one anchor

# 因为一共有9个锚框，但分布在3个尺度下，下面就是将具体的锚框与它所属的尺度对应起来，即找到在每个尺度下的最好的锚框是哪个，将其判断为正例，其他不好的锚框进一步判断为负例还是忽略样例
            # 即对所有的锚框都会做判断，正例的锚框就会计算置信度loss、检测框loss、类别loss，负例只会计算置信度loss，忽略样例则什么loss都不计算
            # 先从9个锚框中重合度最大的锚框开始进行循环
            for anchor_idx in anchor_indices:
                # 根据锚框的索引和每个尺度下拥有的锚框数量，就可以确定锚框所在的尺度
                # 比如如果锚框索引为8，且每个尺度下有3个锚框，那么scale_idx就是2，即第3个尺度，因此scale_idx就是该锚框所属的尺度的索引
                scale_idx = anchor_idx // self.num_anchors_per_scale
                # 上述锚框索引为8，指的是该锚框在所有锚框中的索引，下面就是计算该锚框在该尺度下的索引，即anchor_on_scale就是2，也就是该尺度下该锚框是第3个
                anchor_on_scale = anchor_idx % self.num_anchors_per_scale
                # 获取该尺度下的grid cell的个数，即格点个数
                S = self.S[scale_idx]
                # 提醒一下：边界框的x和y坐标是其中心相对于整张图像的位置
                # 下面就是计算边界框属于图像中的哪个格点
                # 比如假设整个图像宽为W，那么边界框绝对位置就在W*x，而每个格点的宽度为W/S，那么在哪个格点就是W*x/(W/S)=S*x
                # 这里需要注意的是x是宽，但在矩阵中是列，即j,
                # 而y是高，在矩阵中是行，即i
                i, j = int(S * y), int(S * x)  # which cell
                # 下面这一行就是对于这一边界框，取出某一尺度下的锚框、格点及置信度（0元素就是代表是一个物体的可能性即置信度）
                # 刚开始anchor_taken都是0，表明在该尺度的该锚框没有被取走或说没有被判断，非0的话又有两种，1是代表是正例，-1代表是忽略样例
                anchor_taken = targets[scale_idx][anchor_on_scale, i, j, 0]
                # 如果该尺度下（或称该格点）的该锚框没有被判断，且该尺度上之前没有确定锚框（即还没有出现正例）
                if not anchor_taken and not has_anchor[scale_idx]:
                    # 就把第0个元素，即是一个物体的置信度置为1
                    targets[scale_idx][anchor_on_scale, i, j, 0] = 1
                    # 计算边界框的中心在格点中的相对位置
                    x_cell, y_cell = S * x - j, S * y - i  # both between [0,1]
                    # 计算边界框的宽高相对于格点的大小
                    # 仍然假设整张图像宽为W，边界框的绝对宽度就是W*width，那么它相对于格点的大小就是W*width/(W/S)=width*S
                    width_cell, height_cell = (
                        width * S,
                        height * S,
                    )  # can be greater than 1 since it's relative to cell
                    box_coordinates = torch.tensor(
                        [x_cell, y_cell, width_cell, height_cell]
                    )
                    # 然后把上面的边界框相对于格点的相对位置和相对大小信息都存储到targets相应元素中，与具体的尺度、锚框和格点进行匹配。
                    targets[scale_idx][anchor_on_scale, i, j, 1:5] = box_coordinates
                    # 将物体类别也存储到相应元素
                    targets[scale_idx][anchor_on_scale, i, j, 5] = int(class_label)
                    # 将该尺度下是否确定了锚框置为True，即该锚框为正例
                    has_anchor[scale_idx] = True

# 如果已经出现了正例，但该锚框还没有被判断，即anchor_taken=0
                # 此时再判断该锚框与边界框的IoU是否大于阈值，如果大于阈值，且因为其不是正例，那么就将其置信度标签置为-1，即它为忽略样例，不参与损失计算
                # 这种情况出现在在该尺度上（或称在该格点上），有多个锚框都能与边界框吻合较好，但只取最好的那个
                # 但如果IoU小于阈值，那么其置信度标签仍为0，代表负例，会产生置信度loss，但不会产生其他类型的损失
                elif not anchor_taken and iou_anchors[anchor_idx] > self.ignore_iou_thresh:
                    targets[scale_idx][anchor_on_scale, i, j, 0] = -1  # ignore prediction

return image, tuple(targets)

class YoloLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.mse = nn.MSELoss() # 均方差损失计算
        self.bce = nn.BCEWithLogitsLoss() # 加了Sigmoid的二进制交叉熵损失
        self.entropy = nn.CrossEntropyLoss() # 交叉熵损失
        self.sigmoid = nn.Sigmoid() 

# 确定损失函数中的各个权重常数，用来控制不同loss之间的比例
        self.lambda_class = 1 # 类别损失权重常数
        self.lambda_noobj = 10 # 负例损失权重常数
        self.lambda_obj = 1 # 正例损失权重常数
        self.lambda_box = 10 # 检测框损失权重常数

def forward(self, predictions, target, anchors):
        # 判断每个尺度上的格点上是否是物体，即正例还是负例
        # 1为正例，0为负例，-1则为忽略样例
        obj = target[..., 0] == 1  # in paper this is Iobj_i
        noobj = target[..., 0] == 0  # in paper this is Inoobj_i

# ======================= #
        #   负例造成的置信度损失    #
        # ======================= #

# 使用二进制交叉熵计算损失
        no_object_loss = self.bce(
            # 取出置信度数值，即第0个元素
            # 这里使用0:1的形式，而不是直接使用0来取得元素，是为了保持维度不变
            # [noobj] 是使用了numpy的布尔索引，从而取出那些负例
            (predictions[..., 0:1][noobj]), (target[..., 0:1][noobj]),
        )

# ==================== #
        #   正例造成的置信度损失 #
        # ==================== #

# 这个地方正例损失按上面的解析应该使用一个简单的bce即可，同时置信度标签在yolov3中是1和0二分类，而这里原作者使用的是IoU来作为置信度标签，即如下形式：
        object_loss = self.bce(
            (predictions[..., 0:1][obj]), (target[..., 0:1][obj]),
        )

# # 原来的代码中是如下形式，这里先不仔细研究异同了
        # anchors = anchors.reshape(1, 3, 1, 1, 2)
        # box_preds = torch.cat([self.sigmoid(predictions[..., 1:3]), torch.exp(predictions[..., 3:5]) * anchors], dim=-1)
        # ious = intersection_over_union(box_preds[obj], target[..., 1:5][obj]).detach()
        # object_loss = self.mse(self.sigmoid(predictions[..., 0:1][obj]), ious * target[..., 0:1][obj])

# ======================== #
        #  检测框损失               #
        # ======================== #

# 这个地方涉及检测框的解码部分
        # 注意只取出正例造成的损失
        predictions[..., 1:3] = self.sigmoid(predictions[..., 1:3])  # x, y坐标
        target[..., 3:5] = torch.log(
            (1e-16 + target[..., 3:5] / anchors)
        )  # width, height coordinates
        box_loss = self.mse(predictions[..., 1:5][obj], target[..., 1:5][obj])

# ================== #
        #   类别损失   #
        # ================== #

# 计算交叉熵损失，注意只取出正例造成的损失
        class_loss = self.entropy(
            (predictions[..., 5:][obj]), (target[..., 5][obj].long()),
        )

#print("__________________________________")
        #print(self.lambda_box * box_loss)
        #print(self.lambda_obj * object_loss)
        #print(self.lambda_noobj * no_object_loss)
        #print(self.lambda_class * class_loss)
        #print("\n")

return (
            self.lambda_box * box_loss
            + self.lambda_obj * object_loss
            + self.lambda_noobj * no_object_loss
            + self.lambda_class * class_loss
        )

DATASET = 'PASCAL_VOC'
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
# seed_everything()  # If you want deterministic behavior
NUM_WORKERS = 4
BATCH_SIZE = 32
IMAGE_SIZE = 416
NUM_CLASSES = 20 #类别数
LEARNING_RATE = 1e-5
WEIGHT_DECAY = 1e-4
NUM_EPOCHS = 100
CONF_THRESHOLD = 0.05
MAP_IOU_THRESH = 0.5
NMS_IOU_THRESH = 0.45
S = [IMAGE_SIZE // 32, IMAGE_SIZE // 16, IMAGE_SIZE // 8]
PIN_MEMORY = True
LOAD_MODEL = True
SAVE_MODEL = True
CHECKPOINT_FILE = "checkpoint.pth.tar"
IMG_DIR = DATASET + "/images/"
LABEL_DIR = DATASET + "/labels/"

# 通过在训练集上kmeans聚类得到的锚框的大小
ANCHORS = [
    [(0.28, 0.22), (0.38, 0.48), (0.9, 0.78)],
    [(0.07, 0.15), (0.15, 0.11), (0.14, 0.29)],
    [(0.02, 0.03), (0.04, 0.07), (0.08, 0.06)],
]  # Note these have been rescaled to be between [0, 1]

def train_fn(train_loader, model, optimizer, loss_fn, scaler, scaled_anchors):
    # 显示进度条
    loop = tqdm(train_loader, leave=True)
    losses = []

for batch_idx, (x, y) in enumerate(loop):
        x = x.to(config.DEVICE)
        # 三个不同尺度下的目标
        y0, y1, y2 = (
            y[0].to(config.DEVICE),
            y[1].to(config.DEVICE),
            y[2].to(config.DEVICE),
        )

with torch.cuda.amp.autocast():
            out = model(x)
            # 前面的损失函数需要在三个尺度下都要计算一遍
            loss = (
                loss_fn(out[0], y0, scaled_anchors[0])
                + loss_fn(out[1], y1, scaled_anchors[1])
                + loss_fn(out[2], y2, scaled_anchors[2])
            )

losses.append(loss.item())
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

# update progress bar
        mean_loss = sum(losses) / len(losses)
        loop.set_postfix(loss=mean_loss)

def main():
    # 定义模型
    model = YOLOv3(num_classes=config.NUM_CLASSES).to(config.DEVICE)
    # 优化器
    optimizer = optim.Adam(
        model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY
    )
    # 损失函数
    loss_fn = YoloLoss()
    scaler = torch.cuda.amp.GradScaler()

# 数据加载器
    train_loader, test_loader, train_eval_loader = get_loaders(
        train_csv_path=config.DATASET + "/train.csv", test_csv_path=config.DATASET + "/test.csv"
    )

# 可以加载已训练好的模型
    if config.LOAD_MODEL:
        load_checkpoint(
            config.CHECKPOINT_FILE, model, optimizer, config.LEARNING_RATE
        )

scaled_anchors = (
        torch.tensor(config.ANCHORS)
        * torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
    ).to(config.DEVICE)

# 开始迭代训练
    for epoch in range(config.NUM_EPOCHS):
        #plot_couple_examples(model, test_loader, 0.6, 0.5, scaled_anchors)
        train_fn(train_loader, model, optimizer, loss_fn, scaler, scaled_anchors)

#if config.SAVE_MODEL:
        #    save_checkpoint(model, optimizer, filename=f"checkpoint.pth.tar")

#print(f"Currently epoch {epoch}")
        #print("On Train Eval loader:")
        #print("On Train loader:")
        #check_class_accuracy(model, train_loader, threshold=config.CONF_THRESHOLD)

if epoch > 0 and epoch % 3 == 0:
            check_class_accuracy(model, test_loader, threshold=config.CONF_THRESHOLD)

# 得到预测框和真实的边界框的对比
            # 因为对于一张图像，在三个尺度上会有多个预测框与之吻合挺好，这里使用了NMS非极大值抑制来选择出最好的一个预测框
            pred_boxes, true_boxes = get_evaluation_bboxes(
                test_loader,
                model,
                iou_threshold=config.NMS_IOU_THRESH,
                anchors=config.ANCHORS,
                threshold=config.CONF_THRESHOLD,
            )
            # 计算mAP
            mapval = mean_average_precision(
                pred_boxes,
                true_boxes,
                iou_threshold=config.MAP_IOU_THRESH,
                box_format="midpoint",
                num_classes=config.NUM_CLASSES,
            )
            print(f"MAP: {mapval.item()}")
            model.train()