# 根据yolov7改编:https://github.com/WongKinYiu/yolov7 import torch from model.layer import cbs, elan, elan_h, mp, sppcspc, concat, head class yolov7(torch.nn.Module): def __init__(self, args): super().__init__() dim_dict = {'n': 8, 's': 16, 'm': 32, 'l': 64} n_dict = {'n': 1, 's': 1, 'm': 2, 'l': 3} dim = dim_dict[args.model_type] n = n_dict[args.model_type] input_size = args.input_size stride = (8, 16, 32) self.output_size = [int(input_size // i) for i in stride] # 每个输出层的尺寸，如(80,40,20) self.output_class = args.output_class # 网络结构 if not args.prune: # 正常版本 self.l0 = cbs(3, dim, 3, 1) self.l1 = cbs(dim, 2 * dim, 3, 2) # input_size/2 self.l2 = cbs(2 * dim, 2 * dim, 3, 1) self.l3 = cbs(2 * dim, 4 * dim, 3, 2) # input_size/4 # ---------- # self.l4 = elan(4 * dim, 8 * dim, n) self.l5 = mp(8 * dim, 8 * dim) # input_size/8 self.l6 = elan(8 * dim, 16 * dim, n) self.l7 = mp(16 * dim, 16 * dim) # input_size/16 self.l8 = elan(16 * dim, 32 * dim, n) self.l9 = mp(32 * dim, 32 * dim) # input_size/32 self.l10 = elan(32 * dim, 32 * dim, n) self.l11 = sppcspc(32 * dim, 16 * dim) self.l12 = cbs(16 * dim, 8 * dim, 1, 1) # ---------- # self.l13 = torch.nn.Upsample(scale_factor=2) # input_size/16 self.l8_add = cbs(32 * dim, 8 * dim, 1, 1) self.l14 = concat(1) self.l15 = elan_h(16 * dim, 8 * dim) self.l16 = cbs(8 * dim, 4 * dim, 1, 1) # ---------- # self.l17 = torch.nn.Upsample(scale_factor=2) # input_size/8 self.l6_add = cbs(16 * dim, 4 * dim, 1, 1) self.l18 = concat(1) self.l19 = elan_h(8 * dim, 4 * dim) # 接output0 # ---------- # self.l20 = mp(4 * dim, 8 * dim) self.l21 = concat(1) self.l22 = elan_h(16 * dim, 8 * dim) # 接output1 # ---------- # self.l23 = mp(8 * dim, 16 * dim) self.l24 = concat(1) self.l25 = elan_h(32 * dim, 16 * dim) # 接output2 # ---------- # self.output0 = head(4 * dim, self.output_size[0], self.output_class) self.output1 = head(8 * dim, self.output_size[1], self.output_class) self.output2 = head(16 * dim, self.output_size[2], self.output_class) else: # 剪枝版本 config = args.prune_num self.l0 = cbs(3, config[0], 1, 1) self.l1 = cbs(config[0], config[1], 3, 2) # input_size/2 self.l2 = cbs(config[1], config[2], 1, 1) self.l3 = cbs(config[2], config[3], 3, 2) # input_size/4 # ---------- # self.l4 = elan(config[3], None, n, config[4:7 + 2 * n]) self.l5 = mp(config[6 + 2 * n], None, config[7 + 2 * n:10 + 2 * n]) # input_size/8 self.l6 = elan(config[7 + 2 * n] + config[9 + 2 * n], None, n, config[10 + 2 * n:13 + 4 * n]) self.l7 = mp(config[12 + 4 * n], None, config[13 + 4 * n:16 + 4 * n]) # input_size/16 self.l8 = elan(config[13 + 4 * n] + config[15 + 4 * n], None, n, config[16 + 4 * n:19 + 6 * n]) self.l9 = mp(config[18 + 6 * n], None, config[19 + 6 * n:22 + 6 * n]) # input_size/32 self.l10 = elan(config[19 + 6 * n] + config[21 + 6 * n], None, n, config[22 + 6 * n:25 + 8 * n]) self.l11 = sppcspc(config[24 + 8 * n], None, config[25 + 8 * n:32 + 8 * n]) self.l12 = cbs(config[31 + 8 * n], config[32 + 8 * n], 1, 1) # ---------- # self.l13 = torch.nn.Upsample(scale_factor=2) # input_size/16 self.l8_add = cbs(config[18 + 6 * n], config[33 + 8 * n], 1, 1) self.l14 = concat(1) self.l15 = elan_h(config[32 + 8 * n] + config[33 + 8 * n], None, config[34 + 8 * n:41 + 8 * n]) self.l16 = cbs(config[40 + 8 * n], config[41 + 8 * n], 1, 1) # ---------- # self.l17 = torch.nn.Upsample(scale_factor=2) # input_size/8 self.l6_add = cbs(config[12 + 4 * n], config[42 + 8 * n], 1, 1) self.l18 = concat(1) self.l19 = elan_h(config[41 + 8 * n] + config[42 + 8 * n], None, config[43 + 8 * n:50 + 8 * n]) # 接output0 # ---------- # self.l20 = mp(config[49 + 8 * n], None, config[50 + 8 * n:53 + 8 * n]) self.l21 = concat(1) self.l22 = elan_h(config[40 + 8 * n] + config[50 + 8 * n] + config[52 + 8 * n], None, config[53 + 8 * n:60 + 8 * n]) # 接output1 # ---------- # self.l23 = mp(config[59 + 8 * n], None, config[60 + 8 * n:63 + 8 * n]) self.l24 = concat(1) self.l25 = elan_h(config[31 + 8 * n] + config[60 + 8 * n] + config[62 + 8 * n], None, config[63 + 8 * n:70 + 8 * n]) # 接output2 # ---------- # self.output0 = head(config[49 + 8 * n], self.output_size[0], self.output_class) self.output1 = head(config[59 + 8 * n], self.output_size[1], self.output_class) self.output2 = head(config[69 + 8 * n], self.output_size[2], self.output_class) def forward(self, x): x = self.l0(x) x = self.l1(x) x = self.l2(x) x = self.l3(x) x = self.l4(x) x = self.l5(x) l6 = self.l6(x) x = self.l7(l6) l8 = self.l8(x) x = self.l9(l8) x = self.l10(x) l11 = self.l11(x) x = self.l12(l11) x = self.l13(x) l8_add = self.l8_add(l8) x = self.l14([x, l8_add]) l15 = self.l15(x) x = self.l16(l15) x = self.l17(x) l6_add = self.l6_add(l6) x = self.l18([x, l6_add]) x = self.l19(x) output0 = self.output0(x) x = self.l20(x) x = self.l21([x, l15]) x = self.l22(x) output1 = self.output1(x) x = self.l23(x) x = self.l24([x, l11]) x = self.l25(x) output2 = self.output2(x) return [output0, output1, output2] if __name__ == '__main__': import argparse parser = argparse.ArgumentParser(description='') parser.add_argument('--prune', default=False, type=bool) parser.add_argument('--model_type', default='n', type=str) parser.add_argument('--input_size', default=640, type=int) parser.add_argument('--output_class', default=1, type=int) args = parser.parse_args() model = yolov7(args) tensor = torch.rand(2, 3, args.input_size, args.input_size, dtype=torch.float32) pred = model(tensor) print(model) print(pred[0].shape, pred[1].shape, pred[2].shape)