diff --git a/export_my_gra.sh b/export_my_gra.sh index 953f51e..ed354d3 100644 --- a/export_my_gra.sh +++ b/export_my_gra.sh @@ -1,6 +1,6 @@ python export.py \ - --weights /data6/zhangw/ZW_gra_project/yolov5_result_pretrained/G0_cub100_yolov5s6_raw_e300/weights/best.pt \ - --img 640 640 \ + --weights /data6/zhangw/yolov5/yolov5s6.pt \ + --img 640 \ --batch 1 \ --device cpu \ --include onnx \ diff --git a/models/yolo.py b/models/yolo.py index 13498ac..4ffefd9 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -126,6 +126,90 @@ class Detect(nn.Module): anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape) return grid, anchor_grid +# 新增方法 +class DecoupledDetect(nn.Module): + """YOLOv5 decoupled detect head: anchor-based, training/inference interface kept the same as Detect.""" + stride = None + dynamic = False + export = False + + def __init__(self, nc=80, anchors=(), ch=(), inplace=True): + super().__init__() + self.nc = nc # number of classes + self.no = nc + 5 # outputs per anchor = [x, y, w, h, obj] + cls + self.nl = len(anchors) # number of detection layers + self.na = len(anchors[0]) // 2 # number of anchors + self.grid = [torch.empty(0) for _ in range(self.nl)] + self.anchor_grid = [torch.empty(0) for _ in range(self.nl)] + self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2)) + self.inplace = inplace + + # hidden channels for each scale + c_ = [max(128, x // 2) for x in ch] + + # stem + self.stems = nn.ModuleList(Conv(x, c, 1, 1) for x, c in zip(ch, c_)) + + # classification branch + self.cls_convs = nn.ModuleList( + nn.Sequential( + Conv(c, c, 3, 1), + Conv(c, c, 3, 1), + ) for c in c_ + ) + + # regression/objectness branch + self.reg_convs = nn.ModuleList( + nn.Sequential( + Conv(c, c, 3, 1), + Conv(c, c, 3, 1), + ) for c in c_ + ) + + # prediction layers + self.cls_preds = nn.ModuleList(nn.Conv2d(c, self.na * self.nc, 1) for c in c_) + self.reg_preds = nn.ModuleList(nn.Conv2d(c, self.na * 4, 1) for c in c_) + self.obj_preds = nn.ModuleList(nn.Conv2d(c, self.na * 1, 1) for c in c_) + + def forward(self, x): + z = [] # inference output + for i in range(self.nl): + xi = self.stems[i](x[i]) + + cls_feat = self.cls_convs[i](xi) + reg_feat = self.reg_convs[i](xi) + + cls_out = self.cls_preds[i](cls_feat) # [bs, na*nc, ny, nx] + reg_out = self.reg_preds[i](reg_feat) # [bs, na*4, ny, nx] + obj_out = self.obj_preds[i](reg_feat) # [bs, na*1, ny, nx] + + # keep original YOLOv5 output order: [xywh, obj, cls] + x[i] = torch.cat((reg_out, obj_out, cls_out), 1) + + bs, _, ny, nx = x[i].shape + x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous() + + if not self.training: + if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]: + self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i) + + xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4) + xy = (xy * 2 + self.grid[i]) * self.stride[i] + wh = (wh * 2) ** 2 * self.anchor_grid[i] + y = torch.cat((xy, wh, conf), 4) + z.append(y.view(bs, self.na * nx * ny, self.no)) + + return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x) + + def _make_grid(self, nx=20, ny=20, i=0, torch_1_10=check_version(torch.__version__, "1.10.0")): + d = self.anchors[i].device + t = self.anchors[i].dtype + shape = 1, self.na, ny, nx, 2 + y, x = torch.arange(ny, device=d, dtype=t), torch.arange(nx, device=d, dtype=t) + yv, xv = torch.meshgrid(y, x, indexing="ij") if torch_1_10 else torch.meshgrid(y, x) + grid = torch.stack((xv, yv), 2).expand(shape) - 0.5 + anchor_grid = (self.anchors[i] * self.stride[i]).view((1, self.na, 1, 1, 2)).expand(shape) + return grid, anchor_grid class Segment(Detect): """YOLOv5 Segment head for segmentation models, extending Detect with mask and prototype layers.""" @@ -207,7 +291,8 @@ class BaseModel(nn.Module): """ self = super()._apply(fn) m = self.model[-1] # Detect() - if isinstance(m, (Detect, Segment)): + # if isinstance(m, (Detect, Segment)): + if isinstance(m, (Detect, DecoupledDetect, Segment)): m.stride = fn(m.stride) m.grid = list(map(fn, m.grid)) if isinstance(m.anchor_grid, list): @@ -244,7 +329,8 @@ class DetectionModel(BaseModel): # Build strides, anchors m = self.model[-1] # Detect() - if isinstance(m, (Detect, Segment)): + # if isinstance(m, (Detect, Segment)): + if isinstance(m, (Detect, DecoupledDetect, Segment)): def _forward(x): """Passes the input 'x' through the model and returns the processed output.""" @@ -314,26 +400,51 @@ class DetectionModel(BaseModel): y[-1] = y[-1][:, i:] # small return y - def _initialize_biases(self, cf=None): - """ - Initializes biases for YOLOv5's Detect() module, optionally using class frequencies (cf). + # def _initialize_biases(self, cf=None): + # """ + # Initializes biases for YOLOv5's Detect() module, optionally using class frequencies (cf). - For details see https://arxiv.org/abs/1708.02002 section 3.3. - """ + # For details see https://arxiv.org/abs/1708.02002 section 3.3. + # """ + # # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. + # m = self.model[-1] # Detect() module + # for mi, s in zip(m.m, m.stride): # from + # b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) + # b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) + # b.data[:, 5 : 5 + m.nc] += ( + # math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum()) + # ) # cls + # mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) + def _initialize_biases(self, cf=None): + """Initialize biases for Detect() or DecoupledDetect().""" # cf = torch.bincount(torch.tensor(np.concatenate(dataset.labels, 0)[:, 0]).long(), minlength=nc) + 1. m = self.model[-1] # Detect() module - for mi, s in zip(m.m, m.stride): # from - b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) - b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj (8 objects per 640 image) - b.data[:, 5 : 5 + m.nc] += ( - math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum()) - ) # cls - mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) + if isinstance(m, Detect): + for mi, s in zip(m.m, m.stride): + b = mi.bias.view(m.na, -1) # conv.bias(255) to (3,85) + b.data[:, 4] += math.log(8 / (640 / s) ** 2) # obj + b.data[:, 5:5 + m.nc] += ( + math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum()) + ) # cls + mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True) + + elif isinstance(m, DecoupledDetect): + for cls_pred, obj_pred, s in zip(m.cls_preds, m.obj_preds, m.stride): + # objectness bias + b_obj = obj_pred.bias.view(m.na, -1) # [na, 1] + b_obj.data += math.log(8 / (640 / s) ** 2) + obj_pred.bias = torch.nn.Parameter(b_obj.view(-1), requires_grad=True) + + # classification bias + b_cls = cls_pred.bias.view(m.na, -1) # [na, nc] + b_cls.data += ( + math.log(0.6 / (m.nc - 0.99999)) if cf is None else torch.log(cf / cf.sum()) + ) + cls_pred.bias = torch.nn.Parameter(b_cls.view(-1), requires_grad=True) Model = DetectionModel # retain YOLOv5 'Model' class for backwards compatibility - class SegmentationModel(DetectionModel): """YOLOv5 segmentation model for object detection and segmentation tasks with configurable parameters.""" @@ -434,7 +545,8 @@ def parse_model(d, ch): elif m is Concat: c2 = sum(ch[x] for x in f) # TODO: channel, gw, gd - elif m in {Detect, Segment}: + # elif m in {Detect, Segment}: + elif m in {Detect, DecoupledDetect, Segment}: args.append([ch[x] for x in f]) if isinstance(args[1], int): # number of anchors args[1] = [list(range(args[1] * 2))] * len(f) diff --git a/train_my_gra.sh b/train_my_gra.sh index 774e767..8a7dae6 100644 --- a/train_my_gra.sh +++ b/train_my_gra.sh @@ -1,9 +1,10 @@ -export CUDA_VISIBLE_DEVICES=4,5,6,7 +export CUDA_VISIBLE_DEVICES=0,1,2,3 python /data6/zhangw/yolov5/train.py \ --img 960 \ --batch 16 \ --epochs 300 \ --data /data6/zhangw/ZW_gra_project/CUB_100/cub100.yaml \ --weights /data6/zhangw/yolov5/yolov5s6.pt \ - --project /data6/zhangw/ZW_gra_project/yolov5_result \ - --name G0_cub100_yolov5s6_raw_e300 \ No newline at end of file + --project /data6/zhangw/ZW_gra_project/yolov5_result_pretrained \ + --name G6_cub100_yolov5s6_decoupled_head_e300 \ + --cfg /data6/zhangw/yolov5/models/hub/yolov5s6_decoupled.yaml \ No newline at end of file