yolov5

 

https://github.com/TrojanXu/yolov5-tensorrt

 

https://github.com/ultralytics/yolov5

测试代码：

1060上，batch size 12是ok的

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
    parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    opt = parser.parse_args()
    opt.cfg = glob.glob('./**/' + opt.cfg, recursive=True)[0]  # find file
    device = torch_utils.select_device(opt.device)

    # Create model
    model = Model(opt.cfg).to(device)
    model.eval()
    torch.save(model.state_dict(), f'v2.pth')
    input = torch.randn(12, 3, 640, 640).cuda()
    for i in range(10):
        start = time.time()
        output = model(input)
        print('output.size ', time.time() - start, output[0].size())
        del output

 

yolov5m pytorch代码：

1060上，640*640 batch size14 是ok的

import time
from collections import OrderedDict

import torch
import math
import torch.nn as nn
import torch.nn.functional as F


def DWConv(c1, c2, k=1, s=1, act=True):
    # Depthwise convolution
    return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)


class Conv(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):  # ch_in, ch_out, kernel, stride, groups
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, k // 2, groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.LeakyReLU(0.1, inplace=True) if act else nn.Identity()

    def forward(self, x):
        return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))


class Bottleneck(nn.Module):
    # Standard bottleneck
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


class BottleneckCSP(nn.Module):
    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(BottleneckCSP, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
        self.cv4 = Conv(c2, c2, 1, 1)
        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
        self.act = nn.LeakyReLU(0.1, inplace=True)
        self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])

    def forward(self, x):
        y1 = self.cv3(self.m(self.cv1(x)))
        y2 = self.cv2(x)
        return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))


class SPP(nn.Module):
    # Spatial pyramid pooling layer used in YOLOv3-SPP
    def __init__(self, c1, c2, k=(5, 9, 13)):
        super(SPP, self).__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

    def forward(self, x):
        x = self.cv1(x)
        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))


class Flatten(nn.Module):
    # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions
    def forward(self, x):
        return x.view(x.size(0), -1)


class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1):
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, 1)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))


class Concat(nn.Module):
    # Concatenate a list of tensors along dimension
    def __init__(self, dimension=1):
        super(Concat, self).__init__()
        self.d = dimension

    def forward(self, x):
        return torch.cat(x, self.d)




class Yolov5(nn.Module):
    def __init__(self):
        super(Yolov5, self).__init__()


        self.layers_out_filters = [192, 384, 768]
        self.features = nn.Sequential()

        i=0
        self.features.add_module(str(i),Focus(3,48,3))
        i+=1
        self.features.add_module(str(i), Conv(48,96,3,2))
        i +=1
        self.features.add_module(str(i), Bottleneck(96,96))
        i += 1
        self.features.add_module(str(i), Conv(96, 192, 3, 2))
        i += 1
        self.features.add_module(str(i), BottleneckCSP(192, 192,6))
        i += 1
        self.features.add_module(str(i), Conv(192, 384, 3,2))
        i += 1
        self.features.add_module(str(i), BottleneckCSP(384, 384, 6))
        i += 1
        self.features.add_module(str(i), Conv(384, 768, 3, 2))
        i += 1
        self.features.add_module(str(i), SPP(768, 768,[5,9,13]))
        i += 1
        self.features.add_module(str(i), BottleneckCSP(768, 768, 4))
        i += 1
        self.features.add_module(str(i), BottleneckCSP(768, 768, 2,False))

        # self.linear = nn.Linear(nin_transition_layer, num_classes)

    def forward(self, x):
        output = []
        for i, module in enumerate(self.features):
            x = module(x)
            print(i,x.shape)
            if i in [4,6]:
                output.append(x)

        output.append(x)
        return output


if __name__ == '__main__':
    pelee_net = Yolov5()

    pelee_net = pelee_net.cuda()
    torch.save(pelee_net.state_dict(), f'v2.pth')
    input = torch.randn(12, 3, 640, 640).cuda()
    for i in range(10):
        start=time.time()
        output = pelee_net(input)
        print('output.size ', time.time()-start,output[0].size(),output[1].size(),output[2].size())
        del output