目标检测第三篇：基于SSD的目标检测算法 |
您所在的位置：网站首页 › ssd目标检测网络的大小 › 目标检测第三篇：基于SSD的目标检测算法 |
目标检测第三篇：基于SSD的目标检测算法
|
文章目录
SSD简介网络搭建卷积块下采样块主干网多层特征提起层输出头
数据处理形成训练TXTDatasetDataLoaderAnchors生成先验框匹配先验框位置 offset
损失函数训练代码及参考
SSD简介
SSD，全称Single Shot MultiBox Detector，是Wei Liu在ECCV 2016上提出的一种目标检测算法，截至目前是主要的检测框架之一，相比Faster RCNN有明显的速度优势，相比YOLO又有明显的mAP优势，是在 RCNN 系列和 YOLO 系列之外属于单阶段的另外一系列的奠基之作。 论文链接：https://arxiv.org/pdf/1512.02325.pdf 官方代码：https://github.com/weiliu89/caffe/tree/ssd SSD的主要设计理念：根据不同的特征层设置不同大小的先验框，在不同的特征层上建立检测头和分类头，以满足大、小目标检测的需求。关于SSD的具体的技术细节，网上大神给出的解释很多，这里不再赘述。 网络架构如下： 论文中以 VGG16 为主干网络，替换了 VGG16 5_3 层和后面的部分，换成了 3*3 的卷积，再加上多尺度特征层，来实现多个检测和分类头。 卷积块论文中以 卷积 + 激活函数为一个卷积标准模块，这里实现的时候加上了 BatchNorm，具体实现如下： def conv_blk(in_channels, out_channels, stride=1, padding=1): """卷积块""" return nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride,padding=padding), nn.BatchNorm2d(out_channels), nn.ReLU()) 下采样块论文中以 池化层作为下采样层，来实现特征层的大小减半，这里实现的时候，将 卷积 + 池化层作为一个标准下采样模块，具体实现如下： def down_sample_blk(in_channels, out_channels, ceil_mode = False): """下采样块""" return nn.Sequential(conv_blk(in_channels, out_channels), nn.MaxPool2d(2, ceil_mode=ceil_mode)) 主干网使用上面两个模块，按照 VGG16 网络的基本架构实现 主干网：具体实现代码如下： def backbone(input_shape): """搭建vgg16主干网络""" return nn.Sequential( conv_blk(input_shape[0], 64), down_sample_blk(64, 64), conv_blk(64, 128), down_sample_blk(128, 128), conv_blk(128, 256), conv_blk(256, 256), down_sample_blk(256, 512, ceil_mode=True), conv_blk(512, 512), conv_blk(512, 512), conv_blk(512, 512) # vgg16 4-3层 ) 多层特征提起层主干网仅仅到 vgg16 4-3层， 后面的我们需要加上特征提取层，来实现输出不同特征层的需求，具体实现代码如下： class extra_feature(nn.Module): """根据backbone输出的特征图，生成额外的特征图""" def __init__(self, input_shape) -> None: super().__init__() # (3, 300, 300) -> (512, 38, 38) self.out_layer1 = backbone(input_shape) # (512, 38, 38) -> (512, 19, 19) self.out_layer2 = nn.Sequential( nn.MaxPool2d(2, stride=1, padding=1), conv_blk(512, 512), conv_blk(512, 512), down_sample_blk(512, 521), conv_blk(521, 1024), conv_blk(1024, 1024) ) # (512, 19, 19) -> (256, 10, 10) self.out_layer3 = nn.Sequential( conv_blk(1024, 256), conv_blk(256, 512, stride=2) ) # (256, 10, 10) -> (256, 5, 5) self.out_layer4 = nn.Sequential( conv_blk(512, 128), conv_blk(128, 256, stride=2) ) # (256, 5, 5) -> (256, 3, 3) self.out_layer5 = nn.Sequential( conv_blk(256, 128), conv_blk(128, 256, stride=2) ) # (256, 3, 3) -> (256, 1, 1) self.out_layer6 = nn.Sequential( conv_blk(256, 128), nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=0) ) def forward(self, x): out1 = self.out_layer1(x) out2 = self.out_layer2(out1) out3 = self.out_layer3(out2) out4 = self.out_layer4(out3) out5 = self.out_layer5(out4) out6 = self.out_layer6(out5) return out1, out2, out3, out4, out5, out6经过 extra_feature 层，我们将会得到六层特征层的输出，对应上面图片的六层输出层，特征层的大小依次为：（38, 38）、（19， 19）、（10， 10）、（5， 5）、（3， 3）、（1， 1）。 输出头得到六个输出层之后，我们需要对六层输出的通道数目进行调整，使其为我们需要分类的数目和需要输出的检测框的数目的大小。具体代码如下： def mutiboxhead(num_classes): """搭建SSD多尺度检测头""" num_anchors = [4, 6, 6, 6, 4, 4] # 每个尺度的锚框数量 num_channels = [512, 1024, 512, 256, 256, 256] # 每个尺度的通道数量 cls_predictors = [] # 类别预测器 bbox_predictors = [] # 边界框预测器 for i in range(6): cls_predictors.append(nn.Conv2d(num_channels[i], num_anchors[i] * (num_classes + 1), kernel_size=3, padding=1)) bbox_predictors.append(nn.Conv2d(num_channels[i], num_anchors[i] * 4, kernel_size=3, padding=1)) cls_predictors = nn.ModuleList(cls_predictors) bbox_predictors = nn.ModuleList(bbox_predictors) return cls_predictors, bbox_predictors根据上面构建的所有东西来实现我们的SSD网络，实现代码如下： class SSD(nn.Module): def __init__(self, input_shape=(3, 300, 300), num_classes=20, **kwargs): super(SSD, self).__init__(**kwargs) self.num_classes = num_classes self.input_shape = input_shape self.extra_feature = extra_feature(input_shape) self.cls_predictors, self.bbox_predictors = mutiboxhead(num_classes) def forward(self, x): x = self.extra_feature(x) cls_preds = [] bbox_preds = [] for i, feature in enumerate(x): cls_preds.append(self.cls_predictors[i](feature).permute(0, 2, 3, 1)) # (batch_size, num_classes, h, w) -> (batch_size, h, w, num_classes) bbox_preds.append(self.bbox_predictors[i](feature).permute(0, 2, 3, 1)) # (batch_size, 4, h, w) -> (batch_size, h, w, 4) bbox_preds = torch.cat([pred.reshape(pred.shape[0], -1, 4) for pred in bbox_preds], dim=1) # (batch_size, num_anchors, 4) cls_preds = torch.cat([pred.reshape(pred.shape[0], -1, self.num_classes + 1) for pred in cls_preds], dim=1) # (batch_size, num_anchors, num_classes + 1) return bbox_preds, cls_preds # (batch_size, num_anchors, 4), (batch_size, num_anchors, num_classes + 1)如果对搭建的网络参数或者输出的形状不是很放心的话，可以输出和跟踪一下看看具体的网络的架构。 if __name__ == "__main__": x = torch.randn(1, 3, 300, 300) ssd = SSD((3, 300, 300), 20) cls_preds, bbox_preds = ssd(x) print(cls_preds[0].shape, bbox_preds[0].shape) 数据处理一般情况下，我们使用 COCO 或者 VOC 数据集进行训练，这里以 VOC 数据进行讲解，如何加载数据。 形成训练TXTVOC 的标签为 XML 文件，在网络加载数据的时候，不是很方便，这里将其处理为 TXT 文件，方便我们读取。处理的结果如下： 运行完以上程序，会形成 train.txt, test.txt val.txt，分别对应训练集、验证集和测试集。训练的时候读取对应的txt即可。 Dataset获得完上面的 txt， 我们就可以编写自己的 Dataset 函数，形成自己在网络训练时候的数据加载方式。具体代码如下： # 加载数据集的函数 class SSDDataset(Dataset): """ 定义自己的数据集加载方式 """ def __init__(self, data_lines, transform=None): """ 初始化文件路径和数据增强所需的信息 Args: data (list): 文件路径列表 transform: torchvision.transforms """ self.data_lines = data_lines self.transform = transform def __len__(self): """ 返回数据集的总长度 Returns: length (int): 数据集的总长度 """ return len(self.data_lines) def __getitem__(self, index): """ 通过索引返回一个数据 Args: index (int): 索引 Returns: img (Tensor): 图像数据 label (Tensor): 标签数据 """ data_line = self.data_lines[index].split(",") image = np.array((Image.open(data_line[0]).convert("RGB"))) bboxes = np.array([list(map(int, box.split(' '))) for box in data_line[1:]]) if self.transform is not None: augmentations = self.transform(image=image, bboxes=bboxes) image = augmentations["image"] bboxes = augmentations["bboxes"] return image, bboxes也就是通过输入的 index，来获取 图片 和对应的 标签信息， transform 是是否采用数据增强，数据增强的具体代码如下： train_transforms = A.Compose( [ # 保持图片的比例进行图片的放大 A.LongestMaxSize(max_size=300), # 不保证图片比例进行放大，需要的时候会对图片添加 0 A.PadIfNeeded( min_height=300, min_width=300, border_mode=cv2.BORDER_CONSTANT ), # 进行正则化 A.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010], max_pixel_value=255,), # 转变为 tensor ToTensorV2(), ], # 对检测框进行调整 bbox_params=A.BboxParams(format="pascal_voc", label_fields=[],check_each_transform=False), ) val_transforms = A.Compose( [ # 保持图片的比例进行图片的放大 A.LongestMaxSize(max_size=300), # 不保证图片比例进行放大，需要的时候会对图片添加 0 A.PadIfNeeded( min_height=300, min_width=300, border_mode=cv2.BORDER_CONSTANT ), # 进行正则化 A.Normalize(mean=[0.4914, 0.4822, 0.4465], std=[0.2023, 0.1994, 0.2010], max_pixel_value=255,), # 转变为 tensor ToTensorV2(), ], # 对检测框进行调整 bbox_params=A.BboxParams(format="pascal_voc", label_fields=[], check_each_transform=False), )因为我们使用的格式为 pascal_voc 的格式，所以这里的我们选择的增强方式为 pascal_voc，如果使用的为 coco 的话，换成 coco 即可。 一般情况下，在训练的时候和在验证的时候采用的数据增强手段是不一样的，训练的时候，我们强调数据的多样性，会使用较多的数据增强的手段，为反转、随机剪裁、色域变换等，但是在验证的时候，我们强调当前网络的性能，一般采取较少的数据增强手段。 注意：我们在对图片进行数据增强的时候，要保证检测的一致性。 DataLoader刚刚只是定义了如何获取图片，但是如何设置为可以供 pytorch 使用，需要放到 DataLoader 中， 形成迭代器，来实现不间断的拿取 batch_size 的图片。 具体代码如下： def load_data(root_path, batch_size=8, num_workers=0): # 第一步：加载数据集 with open(os.path.join(root_path, "train.txt"), "r") as f: train_lines = f.readlines() with open(os.path.join(root_path, "train.txt"), "r") as f: val_lines = f.readlines() train_dataset = SSDDataset(train_lines, train_transforms) val_dataset = SSDDataset(val_lines, val_transforms) train_iter = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn) test_iter = DataLoader(val_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn) return train_iter, test_iter def collate_fn(batch): """ 自定义一个函数，将一个batch的数据拼接起来 Args: batch (list): batch数据 Returns: images (Tensor): 图像数据 bboxes (list): 边界框数据 """ images = [] bboxes = [] for image, bbox in batch: images.append(image) bboxes.append(torch.tensor(bbox)) return torch.stack(images, dim=0), bboxes Anchors 生成先验框SSD 网络需要生成在不同的特征层生成先验框，以满足对不同大、小物体检测的需求；具体实现代码如下： # SSD 生成 anchors 的方式, 中心宽高的方式 def generate_anchors(feature_maps_size=[38, 19, 10, 5, 3, 1], image_size=300, steps=[8, 16, 32, 64, 100, 300], min_sizes=[30, 60, 111, 162, 213, 264], max_sizes=[60, 111, 162, 213, 264, 315], aspect_ratios=[[2], [2, 3], [2, 3], [2, 3], [2], [2]], clip=True): mean = [] for k, f in enumerate(feature_maps_size): #[38, 19, 10, 5, 3, 1], for i, j in product(range(f), repeat=2): f_k = image_size / steps[k] # unit center x,y cx = (j + 0.5) / f_k cy = (i + 0.5) / f_k # aspect_ratio: 1 # rel size: min_size s_k = min_sizes[k] / image_size mean += [cx, cy, s_k, s_k] # aspect_ratio: 1 # rel size: sqrt(s_k * s_(k+1)) s_k_prime = sqrt(s_k * (max_sizes[k]/image_size)) mean += [cx, cy, s_k_prime, s_k_prime] # rest of aspect ratios for ar in aspect_ratios[k]: mean += [cx, cy, s_k*sqrt(ar), s_k/sqrt(ar)] mean += [cx, cy, s_k/sqrt(ar), s_k*sqrt(ar)] # back to torch land prior_anchors = torch.Tensor(mean).view(-1, 4) # * image_size ##[8732,4] # 对于超出边界的先验框进行裁剪到边界内 if clip: prior_anchors.clamp_(max=1, min=0) # 返回的时候转化为 corner 的数据格式，并且乘上 图片 对应的尺度获取真实大小 return box_center_to_corner(prior_anchors) * image_size 匹配先验框在训练的时候，我们需要将真实的标签框和先验框进行匹配，来获取训练的正样本。也就是先确保每个标签框有一个先验框负责预测，完成 一对一 的匹配。如果其他的先验框没有被分配，且和其中某个标签框的 iou 也很大，也将对列为正样本，完成 先验框 -> 标签框 的 多对一 的匹配。具体技术细节，请参考：沐神的动手学深度学习的 13.4. 锚框 的章节。 具体实现代码如下： def assign_anchor_to_bbox(ground_truth, anchors, device, iou_threshold=0.5): """ 为锚框分配边界框 Args: ground_truth (Tensor): 边界框，形状为 (n, 4) anchors (Tensor): 锚框，形状为 (k, 4) device (torch.device): 用于分配张量的设备 iou_threshold (float): 阈值 Returns: Tensor: 所分配的边界框，形状为 (k, 4) """ num_anchors, num_gt_boxes = anchors.shape[0], ground_truth.shape[0] # 计算锚框和边界框的交并比 anchors_boxes_iou = box_iou(anchors, ground_truth) # 为每个锚框分配边界框的索引, 形状为 (k, ), 默认为 -1 anchor_to_bbox_map = torch.full((num_anchors,), -1, dtype=torch.long, device=device) # 找到每个 anchors 对应的最大 iou 值和索引 max_ious, indices = torch.max(anchors_boxes_iou, dim=1) # 将交并比大于阈值的锚框分配给边界框 anchor_to_bbox_map[max_ious >= iou_threshold] = indices[max_ious >= iou_threshold] # 找到应该丢弃的行和列，也就是找到每一个 gt 真实值匹配的 anchor 所以对应的行和列 col_discard = torch.full((num_anchors, ), -1) row_discard = torch.full((num_gt_boxes,), -1) # 每个 iou 较大的 anchor只保留一个gt对应 for _ in range(num_gt_boxes): # 找到目前 iou 最大的位置，因为没有指定维度的话，默认将所有维度参与 # anchors_boxes_iou 的维度为 ： [num_anchors, num_gt_boxes] max_idx = torch.argmax(anchors_boxes_iou) # 找到对应的列 box_idx = (max_idx % num_gt_boxes).long() # 找到对应的行 anc_idx = (max_idx / num_gt_boxes).long() # 将对应的 gt 作为 anchor 的真值 anchor_to_bbox_map[anc_idx] = box_idx # 取出对应的行和列 anchors_boxes_iou[:, box_idx] = col_discard anchors_boxes_iou[anc_idx, :] = row_discard return anchor_to_bbox_map 位置 offset网络并不是直接预测标签的中心点和宽高，那样会增加网络预测的难度，我们预测的时候，先验框相对于标签框的偏移量，包括中心点的偏移量和宽高的偏移量。参考 沐神动手学深度的讲解如下; 这些偏移量将作为训练的时候，真正的标签。 损失函数有了网络的输出和真正的偏移量的标签，我们就可以定义我们的损失函数了。 关于 SSD 损失讲解，可以参考大神的解释：链接 简单来说，分为两个部分： 根据上面的损失函数，可以定义我们训练的总的损失函数类：具体代码如下： class MultiBoxLoss(): """根据预测的结果计算损失函数""" def __init__(self, num_classes, device, alpha=1) -> None: super().__init__() self.num_classes = num_classes self.image_size = 300 self.prior_anchors = generate_anchors() self.device = device self.alpha = alpha def __call__(self, net_output, labels): """计算损失函数 Args: net_output (tuple): 网络输出，包含回归预测和分类预测， 形状为 loc_preds:(batch_size, num_anchors, 4), conf_preds: (batch_size, num_anchors, 21) labels (Tensor): 标签，形状为 (n, 5)，其中 n 是所有边界框的数量, 5 表示 (类别, x, y, w, h) Returns: Tensor: 损失值 """ # 第一步：获取网络输出 (batch_size, num_anchors, 4), (batch_size, num_anchors, 21) loc_preds, conf_preds = net_output batch_size, device = len(labels), loc_preds.device num_anchors = self.prior_anchors.shape[0] # self.prior_anchors (num_anchors, 4) # 第二步：为不同大小的特征层生成对应的 anchors 找到对应 target 标号, 并且将其转换为相对于 anchor 的偏移量 bbox_labels, bbox_masks, cls_labels = multibox_target(self.prior_anchors, labels) cls_labels = cls_labels.to(device) bbox_labels = bbox_labels.to(device) bbox_masks = bbox_masks.to(device) num_prior_anchors = bbox_masks[:,:, 0].sum() # (batch_size, 1) # 第三步：计算损失 # 将预测的结果转换为 (batch_size, num_anchors, num_classes + 1)， + 1 为背景类 conf_preds = conf_preds.view(batch_size, num_anchors, self.num_classes + 1) # 计算分类损失 cls_loss = self.conf_loss(conf_preds, cls_labels) # 计算位置损失 loc_loss = self.loc_loss(loc_preds, bbox_labels, bbox_masks) all_loss = (cls_loss + self.alpha * loc_loss) / num_prior_anchors return all_loss def loc_loss(self, loc_preds, bbox_labels, bbox_masks): """计算位置损失 Args: loc_preds (Tensor): 预测的位置结果，形状为 (batch_size, num_anchors, 4) bbox_labels (Tensor): 真实的位置标签，形状为 (batch_size, num_anchors, 4) bbox_masks (Tensor): 真实的位置标签的掩码，形状为 (batch_size, num_anchors, 4) Returns: Tensor: 位置损失 """ # 第一步：计算损失 return F.l1_loss(loc_preds * bbox_masks, bbox_labels * bbox_masks, reduction='sum') def conf_loss(self, conf_preds, cls_labels): """计算分类损失 Args: conf_preds (Tensor): 预测的分类结果，形状为 (batch_size, num_anchors, num_classes) cls_labels (Tensor): 真实的分类标签，形状为 (batch_size, num_anchors, 1) Returns: Tensor: 分类损失 """ # 第一步：计算正样本的损失 # 获取正样本的索引 pos_mask = cls_labels > 1 # 获取正样本的数量 num_pos = pos_mask.long().sum(dim=1, keepdim=True) # (batch_size, 1) # 计算正样本的损失, 交叉熵损失函数的输入为 (N, C) 和 (N, )，其中 N 为样本数量，C 为类别数量，会自动找到(N, C)中的对应的值，计算损失 pos_loss = F.cross_entropy( conf_preds.reshape(-1, conf_preds.shape[-1])[pos_mask.reshape(-1)], cls_labels[pos_mask]) # 第二步：计算负样本的损失 # 获取负样本的索引 neg_mask = cls_labels == 0 # 计算负样本的损失 neg_loss = F.cross_entropy(conf_preds.reshape(-1, conf_preds.shape[-1])[neg_mask.reshape(-1)], cls_labels[neg_mask]) # 第三步：计算损失 loss = pos_loss + neg_loss # 返回损失 return loss #-------------------------------# # 计算预测结果 #-------------------------------# def eval(self, net_output, labels, show_img_flag=False): """计算准确率 Args: net_output (tuple): 网络输出，包含回归预测和分类预测，形状为 (loc_preds (1, num_anchors, 4), conf_preds (1, num_anchors, 21)) labels (Tensor): 标签，形状为 (n, 5)，其中 n 是所有边界框的数量, 5 表示 (类别, x, y, w, h) Returns: Tensor: 损失值 """ # 第一步：获取网络输出 (1, num_anchors, 4), (1, num_anchors, 21) loc_preds, conf_preds = net_output batch_size, device = len(labels), loc_preds.device # self.prior_anchors (num_anchors, 4) num_anchors = self.prior_anchors.shape[0] # 第二步：为不同大小的特征层生成对应的 anchors 找到对应 target 标号, 并且将其转换为相对于 anchor 的偏移量 bbox_labels, bbox_masks, cls_labels = multibox_target(self.prior_anchors, labels) cls_labels = cls_labels.to(device) bbox_labels = bbox_labels.to(device) bbox_masks = bbox_masks.to(device) num_prior_anchors = bbox_masks[:,:, 0].sum() # (1,) # 第三步：计算准确度 # 计算分类准确度 cls_acc = self.cls_eval(conf_preds, cls_labels) / (batch_size * num_anchors) #计算分类损失 cls_loss = self.conf_loss(conf_preds, cls_labels) # / (batch_size * num_anchors) # 计算位置损失 bbox_loss = self.bbox_eval(loc_preds, bbox_labels, bbox_masks) # / (batch_size * num_anchors) # 返回准确度 return cls_acc, cls_loss / num_prior_anchors , bbox_loss / num_prior_anchors def cls_eval(self, cls_preds, cls_labels): # 由于类别预测结果放在最后一维，argmax需要指定最后一维。 return float((cls_preds.argmax(dim=-1).type(cls_labels.dtype) == cls_labels).sum()) def bbox_eval(self, bbox_preds, bbox_labels, bbox_masks): return float((torch.abs((bbox_labels - bbox_preds) * bbox_masks)).sum())其中 multibox_target 函数的作用为 ：为不同大小的特征层生成对应的 anchors 找到对应 target 标号, 并且将其转换为相对于 anchor 的偏移量 def multibox_target(anchors, labels): """ 使用真实边界框标记锚框 Args: anchors (Tensor): 先验框，对应的大小为，(num_anchors, 4) labels (Tensor): 多个图片的真实标签，(batch_size, num_labels, 4) 都是 corner 的形式 Returns: bbox_offset (batch_size, num_boxes, 4) bbox_mask (batch_size, num_boxes, 4) class_labels (batch_size, num_boxes, 1) """ # labels 为多个图片的box的标签合集，维度为 [batch_size, num_gt, 4 + 1]， 其中 1 为类别 batch_size, anchors = len(labels), anchors.squeeze(0) batch_offset, batch_mask, batch_class_labels = [], [], [] device, num_anchors = anchors.device, anchors.shape[0] for i in range(batch_size): # 取出第 i 张图片对应的 box 的 gt label = labels[i] # 找到和标签 box 最接近的真实标签框，作为训练的正样本 anchor_bbox_map = assign_anchor_to_bbox(label[:, :4], anchors, device) # 构造 mask 矩阵，大于零的我们认为找打了相应的gt，扩展到4维的目的是为了后面bbox的四个点的mask bbox_mask = ((anchor_bbox_map >= 0).float().unsqueeze(-1)).repeat(1, 4) # 将类标签和分配的边界框坐标初始化为0 class_labels = torch.zeros(num_anchors, dtype=torch.long, device=device) assigned_bbox = torch.zeros((num_anchors, 4), dtype=torch.float32, device=device) # 使用真实边界框来标记锚框的类别 # 如果一个锚框没有被分配，标记为背景，值为零 indices_true = torch.nonzero(anchor_bbox_map >= 0) # 拿出已经分配好的锚框对应的 gt bbox bbox_idx = anchor_bbox_map[indices_true] # 将类别和对应的标签进行赋值。label + 1 的原因是为了和 0 进行区别。 class_labels[indices_true] = label[bbox_idx, 4].long() + 1 assigned_bbox[indices_true] = label[bbox_idx, :4].float() # 偏移量的转换，得到真实框和对应anchor之间的位置偏移 offset = offset_boxes(anchors, assigned_bbox) * bbox_mask # 记录一下对应的关系 batch_offset.append(offset) batch_mask.append(bbox_mask) batch_class_labels.append(class_labels) bbox_offset = torch.stack(batch_offset) bbox_mask = torch.stack(batch_mask) class_labels = torch.stack(batch_class_labels) return bbox_offset, bbox_mask, class_labels 训练定义相关的优化器、学习率等，我们就可以开始我们的训练了；具体训练代码如下; if __name__ == "__main__": # -------------------超参数-------------------# batch_size = 32 num_workers = 4 input_size = (3, 300, 300) num_classes = 20 learning_rate = 1e-2 epochs = 100 weight_decay = 5e-4 root_path = "./data" # -------------------超参数-------------------# # 设置使用的设备 # device = torch.device("cuda" if torch.cuda.is_available() else "cpu") device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") # 第一步：加载数据集 train_iter, test_iter = load_data(root_path, batch_size, num_workers) # 第二步：定义模型 net = SSD(input_size, num_classes) net.to(device) # summary(net,input_size=(3, 300, 300)) # 第三步：定义损失函数 Loss = MultiBoxLoss(num_classes, device) # 第四步：定义优化器 optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate, momentum=0.9, weight_decay=weight_decay) # 第五步：训练模型 train_day = datetime.now().strftime("%Y_%m_%d_%H_%M_%S") print("#", "-"*30, "训练参数", "-"*30, "#") print("训练日期:{0:>65}".format(train_day)) print("训练设备:{0:>65}".format(str(device))) print("训练数据集大小:{0:>59}".format(len(train_iter) * batch_size)) print("测试数据集大小:{0:>59}".format(len(test_iter) * batch_size)) print("训练批次大小:{0:>61}".format(batch_size)) print("训练轮次:{0:>65}".format(epochs)) print("学习率:{0:>67}".format(learning_rate)) print("权重衰减:{0:>65}".format(weight_decay)) print("#", "-"*30, "训练开始", "-"*30, "#") train_file = "./log/train_log_" + train_day os.makedirs(train_file, exist_ok=True) min_loss = 1e+10 for epoch in range(epochs): net.train() phgr = tqdm(train_iter, total=len(train_iter), desc="Train epoch " + str(epoch + 1) + "/" + str(epochs)) epoch_train_losses = [] for i, (images, labels) in enumerate(phgr): # 将数据转换为 cuda 的 tensor images = images.to(device) # 前向传播 net_output = net(images) # 计算损失 net_loss = Loss(net_output, labels) loss = net_loss.item() epoch_train_losses.append(loss) # 反向传播 optimizer.zero_grad() net_loss.backward() optimizer.step() phgr.set_postfix({"loss": loss}) # 保存最优模型 epoch_train_loss = (sum(epoch_train_losses) / len(epoch_train_losses)) print("epoch_train_loss:", epoch_train_loss, " min_loss:", min_loss) if epoch_train_loss epoch + 1}_testloss_{0:.3f}_trainloss_{0:.3f}.pth".format(epoch_test_loss, epoch_train_loss) 代码及参考完整包含预测的代码地址为： 参考 1、沐神的动手学深度学习 |
本文主要是使用 pytorch 对 SSD 进行简单的实现。包括各个模块的讲解、数据的前、后处理、训练参数解释等。
每一行分别为：图片的路径，x_min y_min x_max y_max label … x_min y_min x_max y_max label 在后面加载数据的时候，直接读取即可，稍微方便一点。 处理的代码如下;
具体代码如下：
前面是置信度损失函数，后面是位置损失函数，中间为权重因子。再详细划分如下： 
根据上面公式，定义损失函数如下：【本文地址】
今日新闻 |
推荐新闻 |