模型中存在大量 raster 如何定位问题和优化

[Mo-yun]

模型中存在大量 raster 占用了30%的推理时间 Print <=20 slowest Op for Convolution, larger than 3.00

Sort by time cost ! Node Type Avg(ms) % Called times Flops Rate ReLU 0.003460 0.102208 2.000000 0.000721 Scale 0.008200 0.242227 3.000000 0.062140 Softmax 0.018510 0.546783 6.000000 0.015580 Interp 0.054040 1.596330 7.000000 0.162436 PReLU 0.087658 2.589398 36.000000 0.558610 UnaryOp 0.137259 4.054626 44.000000 0.501555 Deconvolution 0.209790 6.197169 11.000000 17.742975 Reduction 0.213068 6.294011 77.000000 0.107302 BinaryOp 0.424671 12.544737 240.000000 1.959844 While 0.518691 15.322067 63.000000 1.812641 Convolution 0.740798 21.883081 75.000000 73.405533 Raster 0.967516 28.580290 331.000000 3.679464 total time : 3.385255 ms, total mflops : 42.309219 main, 147, cost time: 371.683014 ms

[Mo-yun]

这个归一化模块每个会增加4次Raster

class LayerNorm(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.register_parameter("weight", nn.Parameter(torch.ones([channels])))
        self.register_parameter("bias", nn.Parameter(torch.zeros([channels])))

    def forward(self, x):
        x_var, x_mean = torch.var_mean(x)
        x = (x - x_mean) / (x_var + 1e-7)**0.5
        x = self.weight * x + self.bias
        return x

[Mo-yun]

整体模型是conformer架构的，大致测试了一下，每个全连接层都会增加2次Raster

[jxt1234]

可以 dump 模型结构看下 layernorm 是否被 fuse 了。

self.weight * x 可以改成 x * self.weight

[Mo-yun]

没有被fuse，我只想在最后一维加weight和bias所以没有用torch的官方实现，似乎无法自动fuse。 我把该模块的weight和bias 用broadcast操作拓展之后再过torch的layernorm可以大幅减少Raster数量

[Mo-yun]

还有一个，我发现Depth-wise conv似乎会被拆分成多个卷积，导致卷积数量和Raster数量大幅增加

[jxt1234]

还有一个，我发现Depth-wise conv似乎会被拆分成多个卷积，导致卷积数量和Raster数量大幅增加

当 group > 1 且与 input channel / output channel 不相等，会出现这种情况

[github-actions[bot]]

Marking as stale. No activity in 60 days.