在youlov5中使用BSConv

BSConv-U源码:

import torch
import bsconv.pytorch

class BSConv(torch.nn.Module):

    def __init__(self, num_classes=1000):
        super().__init__()
        self.features = torch.nn.Sequential(
            bsconv.pytorch.BSConvU(3, 32, kernel_size=3, stride=2, padding=1),
            torch.nn.BatchNorm2d(num_features=32),
            torch.nn.ReLU(inplace=True),
            bsconv.pytorch.BSConvU(32, 64, kernel_size=3, stride=2, padding=1),
            torch.nn.BatchNorm2d(num_features=64),
            torch.nn.ReLU(inplace=True),
            bsconv.pytorch.BSConvU(64, 128, kernel_size=3, stride=2, padding=1),
            torch.nn.BatchNorm2d(num_features=128),
            torch.nn.ReLU(inplace=True),
            bsconv.pytorch.BSConvU(128, 256, kernel_size=3, stride=2, padding=1),
            torch.nn.BatchNorm2d(num_features=256),
            torch.nn.ReLU(inplace=True),
            bsconv.pytorch.BSConvU(256, 512, kernel_size=3, stride=2, padding=1),
            torch.nn.BatchNorm2d(num_features=512),
            torch.nn.ReLU(inplace=True),
        )
        self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = torch.nn.Sequential(
            torch.nn.Linear(512, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

放进common.py中,在yolo.py中添加BSConv,需要对上面的代码做修改
1、需要对初始化修改参数,否则会报错
AttributeError: cannot assign module before Module.init() call

def __init__(self, c1, c2, k=1, s=1, g=1)

2、修改前向传播代码,否则最后的输出的shape为1 * class
3、修改下面语句,将自适应平均池化修改成对应特征图的大小

self.avgpool = torch.nn.AdaptiveAvgPool2d((1, 1))

你可能感兴趣的:(pytorch,深度学习,神经网络)