devbox
黑客灵魂
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【PyTorch】基于YOLO的多目标检测项目（二）

作者：黑客灵魂 | 2024-08-02 23:40:06

赞

踩

【PyTorch】基于YOLO的多目标检测项目（二）

【PyTorch】基于YOLO的多目标检测项目（一）

【PyTorch】基于YOLO的多目标检测项目（二）

YOLO-v3网络由跨距为2的卷积层、跳跃连接层和上采样层组成，没有池化层。网络接收一幅416 * 416的图像作为输入，并提供三个YOLO输出。

目录

准备配置文件

搭建YOLO模型 

搭建PyTorch模块

搭建DarkNet模型

定义损失函数

训练模型

部署模型

准备配置文件

新建一个py文件导入以下代码，命名为myutils.py作为配置文件，辅助构建模型。

  1. import torch
  2. from torch import nn
  3.  
  4.  
  5. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  6.  
  7. def parse_model_config(path2file):
  8.     cfg_file = open(path2file, 'r')
  9.     lines = cfg_file.read().split('\n')
  10.  
  11.     lines = [x for x in lines if x and not x.startswith('#')]
  12.     lines = [x.rstrip().lstrip() for x in lines] 
  13.     
  14.     blocks_list = []
  15.     for line in lines:
  16.         if line.startswith('['): 
  17.             blocks_list.append({})
  18.             blocks_list[-1]['type'] = line[1:-1].rstrip()
  19.         else:
  20.             key, value = line.split("=")
  21.             value = value.strip()
  22.             blocks_list[-1][key.rstrip()] = value.strip()
  23.  
  24.     return blocks_list
  25.  
  26.  
  27. def create_layers(blocks_list):
  28.     hyperparams = blocks_list[0]
  29.     channels_list = [int(hyperparams["channels"])]
  30.     module_list = nn.ModuleList()
  31.     
  32.     for layer_ind, layer_dict in enumerate(blocks_list[1:]):
  33.         modules = nn.Sequential()
  34.         
  35.         if layer_dict["type"] == "convolutional":
  36.             filters = int(layer_dict["filters"])
  37.             kernel_size = int(layer_dict["size"])
  38.             pad = (kernel_size - 1) // 2
  39.             bn=layer_dict.get("batch_normalize",0)    
  40.             
  41.             
  42.             conv2d= nn.Conv2d(
  43.                         in_channels=channels_list[-1],
  44.                         out_channels=filters,
  45.                         kernel_size=kernel_size,
  46.                         stride=int(layer_dict["stride"]),
  47.                         padding=pad,
  48.                         bias=not bn)
  49.             modules.add_module("conv_{0}".format(layer_ind), conv2d)
  50.             
  51.             if bn:
  52.                 bn_layer = nn.BatchNorm2d(filters,momentum=0.9, eps=1e-5)
  53.                 modules.add_module("batch_norm_{0}".format(layer_ind), bn_layer)
  54.                 
  55.                 
  56.             if layer_dict["activation"] == "leaky":
  57.                 activn = nn.LeakyReLU(0.1)
  58.                 modules.add_module("leaky_{0}".format(layer_ind), activn)
  59.                 
  60.         elif layer_dict["type"] == "upsample":
  61.             stride = int(layer_dict["stride"])
  62.             upsample = nn.Upsample(scale_factor = stride)
  63.             modules.add_module("upsample_{}".format(layer_ind), upsample) 
  64.             
  65.  
  66.         elif layer_dict["type"] == "shortcut":
  67.             backwards=int(layer_dict["from"])
  68.             filters = channels_list[1:][backwards]
  69.             modules.add_module("shortcut_{}".format(layer_ind), EmptyLayer())
  70.             
  71.         elif layer_dict["type"] == "route":
  72.             layers = [int(x) for x in layer_dict["layers"].split(",")]
  73.             filters = sum([channels_list[1:][l] for l in layers])
  74.             modules.add_module("route_{}".format(layer_ind), EmptyLayer())
  75.             
  76.         elif layer_dict["type"] == "yolo":
  77.             anchors = [int(a) for a in layer_dict["anchors"].split(",")]
  78.             anchors = [(anchors[i], anchors[i + 1]) for i in range(0, len(anchors), 2)]
  79.  
  80.             mask = [int(m) for m in layer_dict["mask"].split(",")]
  81.             
  82.             anchors = [anchors[i] for i in mask]
  83.             
  84.             num_classes = int(layer_dict["classes"])
  85.             img_size = int(hyperparams["height"])
  86.             
  87.             yolo_layer = YOLOLayer(anchors, num_classes, img_size)
  88.             modules.add_module("yolo_{}".format(layer_ind), yolo_layer)
  89.             
  90.         module_list.append(modules)       
  91.         channels_list.append(filters)
  92.  
  93.     return hyperparams, module_list        
  94.  
  95.  
  96.  
  97. class EmptyLayer(nn.Module):
  98.     def __init__(self):
  99.         super(EmptyLayer, self).__init__()
  100.         
  101.         
  102. class YOLOLayer(nn.Module):
  103.  
  104.     def __init__(self, anchors, num_classes, img_dim=416):
  105.         super(YOLOLayer, self).__init__()
  106.         self.anchors = anchors
  107.         self.num_anchors = len(anchors)
  108.         self.num_classes = num_classes
  109.         self.img_dim = img_dim
  110.         self.grid_size = 0 
  111.         
  112.         
  113.     def forward(self, x_in):
  114.         batch_size = x_in.size(0)
  115.         grid_size = x_in.size(2)
  116.         devide=x_in.device
  117.         
  118.         prediction=x_in.view(batch_size, self.num_anchors, 
  119.                              self.num_classes + 5, grid_size, grid_size)
  120.         prediction=prediction.permute(0, 1, 3, 4, 2)
  121.         prediction=prediction.contiguous()
  122.         
  123.         obj_score = torch.sigmoid(prediction[..., 4]) 
  124.         pred_cls = torch.sigmoid(prediction[..., 5:]) 
  125.         
  126.         if grid_size != self.grid_size:
  127.             self.compute_grid_offsets(grid_size, cuda=x_in.is_cuda)
  128.             
  129.         pred_boxes=self.transform_outputs(prediction) 
  130.         
  131.         output = torch.cat(
  132.             (
  133.                 pred_boxes.view(batch_size, -1, 4),
  134.                 obj_score.view(batch_size, -1, 1),
  135.                 pred_cls.view(batch_size, -1, self.num_classes),
  136.             ), -1,)
  137.         return output        
  138.     
  139.     
  140.         
  141.     def compute_grid_offsets(self, grid_size, cuda=True):
  142.         self.grid_size = grid_size
  143.         self.stride = self.img_dim / self.grid_size
  144.         
  145.         self.grid_x = torch.arange(grid_size, device=device).repeat(1, 1, grid_size, 1 ).type(torch.float32)
  146.         self.grid_y = torch.arange(grid_size, device=device).repeat(1, 1, grid_size, 1).transpose(3, 2).type(torch.float32)
  147.         
  148.         scaled_anchors=[(a_w / self.stride, a_h / self.stride) for a_w, a_h in self.anchors]
  149.         self.scaled_anchors=torch.tensor(scaled_anchors,device=device)
  150.         
  151.         self.anchor_w = self.scaled_anchors[:, 0:1].view((1, self.num_anchors, 1, 1))
  152.         self.anchor_h = self.scaled_anchors[:, 1:2].view((1, self.num_anchors, 1, 1))
  153.         
  154.         
  155.         
  156.     def transform_outputs(self,prediction):
  157.         device=prediction.device
  158.         x = torch.sigmoid(prediction[..., 0]) # Center x
  159.         y = torch.sigmoid(prediction[..., 1]) # Center y
  160.         w = prediction[..., 2] # Width
  161.         h = prediction[..., 3] # Height
  162.  
  163.         pred_boxes = torch.zeros_like(prediction[..., :4]).to(device)
  164.         pred_boxes[..., 0] = x.data + self.grid_x
  165.         pred_boxes[..., 1] = y.data + self.grid_y
  166.         pred_boxes[..., 2] = torch.exp(w.data) * self.anchor_w
  167.         pred_boxes[..., 3] = torch.exp(h.data) * self.anchor_h
  168.         
  169.         return pred_boxes * self.stride

搭建YOLO模型 

解析配置文件，使用parse_model_config助手读取并打印

  1. from myutils import parse_model_config
  2.  
  3. path2config="./config/yolov3.cfg"
  4. blocks_list = parse_model_config(path2config)
  5. blocks_list[:2]

搭建PyTorch模块

基于解析的配置文件创建PyTorch模块，调用 create_layers 辅助函数进行转换并获取 PyTorch 模块的列表

  1. from myutils import create_layers
  2.  
  3. hy_pa, m_l= create_layers(blocks_list)
  4. print(m_l)
  5. print(hy_pa)

搭建DarkNet模型

  1. from torch import nn
  2.  
  3. class Darknet(nn.Module):
  4.     def __init__(self, config_path, img_size=416):
  5.         super(Darknet, self).__init__()
  6.         self.blocks_list = parse_model_config(config_path)
  7.         self.hyperparams, self.module_list = create_layers(self.blocks_list)
  8.         self.img_size = img_size
  9.         
  10.     def forward(self, x):
  11.         img_dim = x.shape[2]
  12.         layer_outputs, yolo_outputs = [], []
  13.         
  14.         for block, module in zip(self.blocks_list[1:], self.module_list):
  15.             if block["type"] in ["convolutional", "upsample", "maxpool"]:
  16.                 x = module(x)        
  17.                 
  18.                 
  19.             elif block["type"] == "shortcut":
  20.                 layer_ind = int(block["from"])
  21.                 x = layer_outputs[-1] + layer_outputs[layer_ind]
  22.             elif block["type"] == "yolo":
  23.                 x= module[0](x)
  24.                 yolo_outputs.append(x)
  25.             elif block["type"] == "route":
  26.                 x = torch.cat([layer_outputs[int(l_i)] 
  27.                                for l_i in block["layers"].split(",")], 1)
  28.             layer_outputs.append(x)
  29.         yolo_out_cat = torch.cat(yolo_outputs, 1)
  30.         return yolo_out_cat, yolo_outputs        
  31.     
  32. model = Darknet(path2config).to(device)
  33. print(model)

  1. # 创建一个随机的dummy_img，大小为1x3x416x416，并将其移动到指定的设备上
  2. dummy_img=torch.rand(1,3,416,416).to(device)
  3. # 在不计算梯度的情况下，执行模型的前向传播
  4. with torch.no_grad():
  5.     # 获取模型的前向传播结果
  6.     dummy_out_cat, dummy_out=model.forward(dummy_img)
  7.     # 打印dummy_out_cat的形状
  8.     print(dummy_out_cat.shape)
  9.     # 打印dummy_out中每个元素的形状
  10.     print(dummy_out[0].shape,dummy_out[1].shape,dummy_out[2].shape)

定义损失函数

YOLO通常使用组合损失函数

  1. def get_loss_batch(output,targets, params_loss, opt=None):
  2.     # 获取损失函数的参数
  3.     ignore_thres=params_loss["ignore_thres"]
  4.     scaled_anchors= params_loss["scaled_anchors"]    
  5.     mse_loss= params_loss["mse_loss"]
  6.     bce_loss= params_loss["bce_loss"]
  7.     
  8.     # 获取yolo的参数
  9.     num_yolos=params_loss["num_yolos"]
  10.     num_anchors= params_loss["num_anchors"]
  11.     obj_scale= params_loss["obj_scale"]
  12.     noobj_scale= params_loss["noobj_scale"]
  13.     
  14.     # 初始化损失
  15.     loss=0.0
  16.     for yolo_ind in range(num_yolos):
  17.         # 获取yolo的输出
  18.         yolo_out=output[yolo_ind]
  19.         batch_size, num_bbxs, _=yolo_out.shape
  20.         
  21.         # 获取网格大小
  22.         gz_2=num_bbxs/num_anchors
  23.         grid_size=int(np.sqrt(gz_2))
  24.         
  25.         # 将yolo的输出reshape为(batch_size,num_anchors,grid_size,grid_size,-1)
  26.         yolo_out=yolo_out.view(batch_size,num_anchors,grid_size,grid_size,-1)
  27.         
  28.         # 获取预测的边界框
  29.         pred_boxes=yolo_out[:,:,:,:,:4]
  30.         x,y,w,h= transform_bbox(pred_boxes, scaled_anchors[yolo_ind])
  31.         # 获取预测的置信度
  32.         pred_conf=yolo_out[:,:,:,:,4]
  33.         # 获取预测的类别概率
  34.         pred_cls_prob=yolo_out[:,:,:,:,5:]
  35.         
  36.         # 获取yolo的目标
  37.         yolo_targets = get_yolo_targets({
  38.                                             "pred_cls_prob": pred_cls_prob,
  39.                                             "pred_boxes":pred_boxes,    
  40.                                             "targets": targets,    
  41.                                             "anchors": scaled_anchors[yolo_ind],    
  42.                                             "ignore_thres": ignore_thres,
  43.                                         }) 
  44.         
  45.         # 获取目标掩码
  46.         obj_mask=yolo_targets["obj_mask"]        
  47.         noobj_mask=yolo_targets["noobj_mask"]            
  48.         # 获取目标的x,y,w,h
  49.         tx=yolo_targets["tx"]                
  50.         ty=yolo_targets["ty"]                    
  51.         tw=yolo_targets["tw"]                        
  52.         th=yolo_targets["th"]                            
  53.         # 获取目标的类别
  54.         tcls=yolo_targets["tcls"]                                
  55.         # 获取目标的置信度
  56.         t_conf=yolo_targets["t_conf"]
  57.         
  58.         # 计算x,y,w,h的损失
  59.         loss_x = mse_loss(x[obj_mask], tx[obj_mask])
  60.         loss_y = mse_loss(y[obj_mask], ty[obj_mask])
  61.         loss_w = mse_loss(w[obj_mask], tw[obj_mask])
  62.         loss_h = mse_loss(h[obj_mask], th[obj_mask])
  63.         
  64.         # 计算置信度的损失
  65.         loss_conf_obj = bce_loss(pred_conf[obj_mask], t_conf[obj_mask])
  66.         loss_conf_noobj = bce_loss(pred_conf[noobj_mask], t_conf[noobj_mask])
  67.         loss_conf = obj_scale * loss_conf_obj + noobj_scale * loss_conf_noobj
  68.         # 计算类别的损失
  69.         loss_cls = bce_loss(pred_cls_prob[obj_mask], tcls[obj_mask])
  70.         # 累加损失
  71.         loss += loss_x + loss_y + loss_w + loss_h + loss_conf + loss_cls
  72.         
  73.     # 如果有优化器，则进行反向传播和优化
  74.     if opt is not None:
  75.         opt.zero_grad()
  76.         loss.backward()
  77.         opt.step()
  78.         
  79.     # 返回损失
  80.     return loss.item()        
  1. def transform_bbox(bbox, anchors):
  2.     # 将bbox的x、y、w、h分别赋值给x、y、w、h
  3.     x=bbox[:,:,:,:,0]
  4.     y=bbox[:,:,:,:,1]
  5.     w=bbox[:,:,:,:,2]
  6.     h=bbox[:,:,:,:,3]
  7.     # 将anchors的w、h分别赋值给anchor_w、anchor_h
  8.     anchor_w = anchors[:, 0].view((1, 3, 1, 1))
  9.     anchor_h = anchors[:, 1].view((1, 3, 1, 1))       
  10.     
  11.     # 将x、y分别减去其向下取整的值
  12.     x=x-x.floor()
  13.     y=y-y.floor()
  14.     # 将w、h分别除以anchor_w、anchor_h，并取对数
  15.     w= torch.log(w / anchor_w + 1e-16)
  16.     h= torch.log(h / anchor_h + 1e-16)
  17.     return x, y, w, h
  18.  
  19. def get_yolo_targets(params):
  20.     # 获取预测框、预测类别概率、目标、锚点、忽略阈值
  21.     pred_boxes=params["pred_boxes"]
  22.     pred_cls_prob=params["pred_cls_prob"]
  23.     target=params["targets"]
  24.     anchors=params["anchors"] 
  25.     ignore_thres=params["ignore_thres"] 
  26.  
  27.     # 获取批量大小、锚点数量、网格大小、类别数量
  28.     batch_size = pred_boxes.size(0)
  29.     num_anchors = pred_boxes.size(1)
  30.     grid_size = pred_boxes.size(2)
  31.     num_cls = pred_cls_prob.size(-1)
  32.     
  33.     
  34.     # 定义目标张量的形状
  35.     sizeT=batch_size, num_anchors, grid_size, grid_size
  36.     # 定义目标张量，用于存储目标框的掩码
  37.     obj_mask = torch.zeros(sizeT,device=device,dtype=torch.uint8)
  38.     # 定义目标张量，用于存储非目标框的掩码
  39.     noobj_mask = torch.ones(sizeT,device=device,dtype=torch.uint8)
  40.     # 定义目标张量，用于存储目标框的x坐标
  41.     tx = torch.zeros(sizeT, device=device, dtype=torch.float32)
  42.     # 定义目标张量，用于存储目标框的y坐标
  43.     ty= torch.zeros(sizeT, device=device, dtype=torch.float32)
  44.     # 定义目标张量，用于存储目标框的宽度
  45.     tw= torch.zeros(sizeT, device=device, dtype=torch.float32)
  46.     # 定义目标张量，用于存储目标框的高度
  47.     th= torch.zeros(sizeT, device=device, dtype=torch.float32)
  48.     
  49.     # 定义目标张量的形状
  50.     sizeT=batch_size, num_anchors, grid_size, grid_size, num_cls
  51.     # 定义目标张量，用于存储目标类别
  52.     tcls= torch.zeros(sizeT, device=device, dtype=torch.float32)
  53.     
  54.     # 将目标框的坐标乘以网格大小
  55.     target_bboxes = target[:, 2:] * grid_size
  56.     # 获取目标框的xy坐标
  57.     t_xy = target_bboxes[:, :2]
  58.     # 获取目标框的wh坐标
  59.     t_wh = target_bboxes[:, 2:]
  60.     # 获取目标框的x坐标
  61.     t_x, t_y = t_xy.t()
  62.     # 获取目标框的宽度
  63.     t_w, t_h = t_wh.t()
  64.  
  65.     # 获取目标框的网格坐标
  66.     grid_i, grid_j = t_xy.long().t()
  67.     
  68.     # 计算每个锚点与目标框的iou
  69.     iou_with_anchors=[get_iou_WH(anchor, t_wh) for anchor in anchors]
  70.     # 将iou转换为张量
  71.     iou_with_anchors = torch.stack(iou_with_anchors)
  72.     # 获取iou最大的锚点索引
  73.     best_iou_wa, best_anchor_ind = iou_with_anchors.max(0)
  74.     
  75.     # 获取目标框的batch索引和类别标签
  76.     batch_inds, target_labels = target[:, :2].long().t()
  77.     # 将目标框的掩码设置为1
  78.     obj_mask[batch_inds, best_anchor_ind, grid_j, grid_i] = 1
  79.     # 将非目标框的掩码设置为0
  80.     noobj_mask[batch_inds, best_anchor_ind, grid_j, grid_i] = 0
  81.  
  82.     # 将大于忽略阈值的iou对应的非目标框掩码设置为0
  83.     for ind, iou_wa in enumerate(iou_with_anchors.t()):
  84.         noobj_mask[batch_inds[ind], iou_wa > ignore_thres, grid_j[ind], grid_i[ind]] = 0
  85.         
  86.         
  87.     # 将目标框的x坐标减去网格的整数部分
  88.     tx[batch_inds, best_anchor_ind, grid_j, grid_i] = t_x - t_x.floor()
  89.     # 将目标框的y坐标减去网格的整数部分
  90.     ty[batch_inds, best_anchor_ind, grid_j, grid_i] = t_y - t_y.floor()
  91.     
  92.  
  93.     # 获取最佳锚点的宽度
  94.     anchor_w=anchors[best_anchor_ind][:, 0]
  95.     # 将目标框的宽度除以锚点的宽度，并取对数
  96.     tw[batch_inds, best_anchor_ind, grid_j, grid_i] = torch.log(t_w / anchor_w + 1e-16)
  97.     
  98.     # 获取最佳锚点的高度
  99.     anchor_h=anchors[best_anchor_ind][:, 1]
  100.     # 将目标框的高度除以锚点的高度，并取对数
  101.     th[batch_inds, best_anchor_ind, grid_j, grid_i] = torch.log(t_h / anchor_h + 1e-16)
  102.     
  103.     # 将目标类别设置为1
  104.     tcls[batch_inds, best_anchor_ind, grid_j, grid_i, target_labels] = 1
  105.     
  106.     # 返回目标张量
  107.     output={
  108.         "obj_mask" : obj_mask,
  109.         "noobj_mask" : noobj_mask,
  110.         "tx": tx,
  111.         "ty": ty,
  112.         "tw": tw,
  113.         "th": th,
  114.         "tcls": tcls,
  115.         "t_conf": obj_mask.float(),
  116.     }
  117.     return output    
  118.  
  119. def get_iou_WH(wh1, wh2):
  120.     # 将wh2转置
  121.     wh2 = wh2.t()
  122.     # 获取wh1的宽度和高度
  123.     w1, h1 = wh1[0], wh1[1]
  124.     # 获取wh2的宽度和高度
  125.     w2, h2 = wh2[0], wh2[1]
  126.     # 计算交集面积
  127.     inter_area = torch.min(w1, w2) * torch.min(h1, h2)
  128.     # 计算并集面积
  129.     union_area = (w1 * h1 + 1e-16) + w2 * h2 - inter_area
  130.     # 返回交集面积与并集面积的比值
  131.     return inter_area / union_area

训练模型

在训练数据上训练模型，并在验证数据上对其进行评估，训练过程遵循标准的随机梯度下降（SGD）。

  1. def loss_epoch(model,params_loss,dataset_dl,sanity_check=False,opt=None):
  2.     running_loss=0.0
  3.     len_data=len(dataset_dl.dataset)
  4.     running_metrics= {}
  5.     
  6.     for xb, yb,_ in dataset_dl:
  7.         yb=yb.to(device)
  8.         _,output=model(xb.to(device))
  9.         loss_b=get_loss_batch(output,yb, params_loss,opt)
  10.         running_loss+=loss_b
  11.         if sanity_check is True:
  12.             break 
  13.     loss=running_loss/float(len_data)
  14.     return loss
  15.  
  16. import copy
  17. def train_val(model, params):
  18.     num_epochs=params["num_epochs"]
  19.     params_loss=params["params_loss"]
  20.     opt=params["optimizer"]
  21.     train_dl=params["train_dl"]
  22.     val_dl=params["val_dl"]
  23.     sanity_check=params["sanity_check"]
  24.     lr_scheduler=params["lr_scheduler"]
  25.     path2weights=params["path2weights"]
  26.     
  27.     
  28.     loss_history={
  29.         "train": [],
  30.         "val": [],
  31.     }
  32.     best_model_wts = copy.deepcopy(model.state_dict())
  33.     best_loss=float('inf') 
  34.     
  35.     for epoch in range(num_epochs):
  36.         current_lr=get_lr(opt)
  37.         print('Epoch {}/{}, current lr={}'.format(epoch, num_epochs - 1, current_lr)) 
  38.         model.train()
  39.         train_loss=loss_epoch(model,params_loss,train_dl,sanity_check,opt)
  40.         loss_history["train"].append(train_loss)
  41.         print("train loss: %.6f" %(train_loss))    
  42.         
  43.         model.eval()
  44.         with torch.no_grad():
  45.             val_loss=loss_epoch(model,params_loss,val_dl,sanity_check)
  46.         loss_history["val"].append(val_loss)
  47.         print("val loss: %.6f" %(val_loss))
  48.         
  49.         
  50.         if val_loss < best_loss:
  51.             best_loss = val_loss
  52.             best_model_wts = copy.deepcopy(model.state_dict())
  53.             torch.save(model.state_dict(), path2weights)
  54.             print("Copied best model weights!")
  55.             
  56.         lr_scheduler.step(val_loss)
  57.         if current_lr != get_lr(opt):
  58.             print("Loading best model weights!")
  59.             model.load_state_dict(best_model_wts) 
  60.         print("-"*10) 
  61.     model.load_state_dict(best_model_wts)
  62.     return model, loss_history            
  63.  
  64. def get_lr(opt):
  65.     for param_group in opt.param_groups:
  66.         return param_group['lr']
  1. from torch import optim
  2. from torch.optim.lr_scheduler import ReduceLROnPlateau
  3.  
  4. opt = optim.Adam(model.parameters(), lr=1e-3)
  5. lr_scheduler = ReduceLROnPlateau(opt, mode='min',factor=0.5, patience=20,verbose=1)
  6.  
  7. path2models= "./models/"
  8. if not os.path.exists(path2models):
  9.         os.mkdir(path2models)
  10.         
  11. scaled_anchors=[model.module_list[82][0].scaled_anchors,
  12.                 model.module_list[94][0].scaled_anchors,
  13.                 model.module_list[106][0].scaled_anchors]        
  14.  
  15. mse_loss = nn.MSELoss(reduction="sum")
  16. bce_loss = nn.BCELoss(reduction="sum")
  17. params_loss={
  18.     "scaled_anchors" : scaled_anchors,
  19.     "ignore_thres": 0.5,
  20.     "mse_loss": mse_loss,
  21.     "bce_loss": bce_loss,
  22.     "num_yolos": 3,
  23.     "num_anchors": 3,
  24.     "obj_scale": 1,
  25.     "noobj_scale": 100,
  26. } 
  27.  
  28. params_train={
  29.     "num_epochs": 5,
  30.     "optimizer": opt,
  31.     "params_loss": params_loss,
  32.     "train_dl": train_dl,
  33.     "val_dl": val_dl,
  34.     "sanity_check": True,
  35.     "lr_scheduler": lr_scheduler,
  36.     "path2weights": path2models+"weights.pt",
  37. }
  38. model,loss_hist=train_val(model,params_train)

部署模型

将训练后的权重加载到模型中 

  1. path2weights="./models/weights.pt"
  2. model.load_state_dict(torch.load(path2weights))

  1. img,tg,_=coco_val[11]
  2. print(img.shape)
  3. print(tg.shape)
  4. show_img_bbox(img,tg)

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/黑客灵魂/article/detail/920524
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号