基于君正magik-transform-tool对一维时序模型进行量化

由于一维时序数据为形如[1,23]的特征，当batch_size设置为64时，经dataloader读取传入后，input_tensor为[64,23]的特征。由于magik近支持pool2d,所以需特征进行补0然后转换为[64,5,5]的张量。

        x = torch.cat((x, torch.zeros(x.size(0), 2)), dim=1) #[64, 23] -> [64, 25]
        x = x.view(x.size(0), 5, 5) #[64, 5, 5]

而由于采用的panns模型，默认的conv如下，且有四层ConvBlock结构

class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ConvBlock, self).__init__()

        self.conv1 = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=(3, 3),
                               stride=(1,1),
                               padding=(1,1))
        self.conv2 = nn.Conv2d(in_channels=out_channels,
                               out_channels=out_channels,
                               kernel_size=(3,3),
                               stride=(1,1),
                               padding=(1,1))
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)

    def forward(self, x, pool_size=(2,2), pool_type='avg'):
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = F.relu(x)

        if pool_type == 'max':
            x = F.max_pool2d(x, kernel_size=pool_size)
        elif pool_type == 'avg':
            x = F.avg_pool2d(x, kernel_size=pool_size)
        elif pool_type == 'avg+max':
            x = F.avg_pool2d(x, kernel_size=pool_size) + F.max_pool2d(x, kernel_size=pool_size)
        else:
            raise Exception(
                f'Pooling type of {pool_type} is not supported. It must be one of "max", "avg" and "avg+max".')
        return x

 网络结构定义如下：

    def __init__(self, num_class, input_size, dropout=0.1, extract_embedding: bool = True):

        super(PANNS_CNN10, self).__init__()
        self.bn0 = nn.BatchNorm2d(input_size)
        self.conv_block1 = ConvBlock(in_channels=1, out_channels=64)
        self.conv_block2 = ConvBlock(in_channels=64, out_channels=128)
        self.conv_block3 = ConvBlock(in_channels=128, out_channels=256)
        self.conv_block4 = ConvBlock(in_channels=256, out_channels=512)

        self.fc1 = nn.Linear(self.emb_size, self.emb_size) #512
        self.extract_embedding = extract_embedding

        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(self.emb_size, num_class)

    def forward(self, x):
        x = torch.cat((x, torch.zeros(x.size(0), 2)), dim=1) #[64, 23] -> [64, 25]
        x = x.view(x.size(0), 5, 5) #[64, 5, 5]
        x = x.unsqueeze(1) #[64, 1, 5, 5],将3D张量拓展维4D
        x = x.permute([0, 3, 2, 1]) #[64, 5, 5, 1] input_size: 5
        x = self.bn0(x)
        x = x.permute([0, 3, 2, 1])


        x = self.conv_block1(x, pool_size=(2,2), pool_type='avg')
        x = F.dropout(x, p=0.2, training=self.training)

        x = self.conv_block2(x, pool_size=(2,2), pool_type='avg')
        x = F.dropout(x, p=0.2, training=self.training)

        x = self.conv_block3(x, pool_size=(1, 1), pool_type='avg')
        x = F.dropout(x, p=0.2, training=self.training)
        
        x = self.conv_block4(x, pool_size=(1, 1), pool_type='avg')
        x = F.dropout(x, p=0.2, training=self.training)

        x = x.mean(dim=3)
        x = x.max(dim=2)[0] + x.mean(dim=2)

        x = F.dropout(x, p=0.5, training=self.training)
        x = F.relu(self.fc1(x))

        if self.extract_embedding:
            output = F.dropout(x, p=0.5, training=self.training)
        else:
            output = F.sigmoid(self.fc_audioset(x))

        x = self.dropout(output)
        logits = self.fc(x)

        return logits

修改池化参数

且默认的pool_kernel为(2,2)，所以要将conv_block3,conv_block4的pool改为(1,1)

之后导出为pth模型，再转为onnx后进行量化 参考

在前向传播的过程中，concat tensor需要考虑到device，可以用self.training判断现在所在的时期

x = torch.cat((x, torch.zeros(x.size(0), 6).to("cuda")), dim=1) if self.training else torch.cat((x, torch.zeros(x.size(0), 6)), dim=1)