Can‘t get attribute ‘SiLU‘ on ＜module ‘torch.nn.modules.activation

 

在yolov5 4.0版本测试的时候，加载模型报错：

Can't get attribute 'SiLU' on <module 'torch.nn.modules.activation

 

报错代码：

aaa=torch.load(w, map_location=map_location)['model'].float().fuse().eval()

原因：

torch加载的时候，会自动映射到类里面，但是找不到这个类，就会报错：

临时解决方法：

把类SiLU的代码粘贴到activation文件中。

 

class SiLU(nn.Module):  # export-friendly version of nn.SiLU()
    @staticmethod
    def forward(x):
        return x * torch.sigmoid(x)

 

这一版没错过，inplace为true，这样可能会减少内存。