Fantastic, thanks @ilovescience a ton for the link!
I’m going to try and extricate the EfficientNet out so we have a pure ENet codebase but looks like he’s already solved the TF issues that I wasn’t sure of how to translate inito PyTorch…so this makes it 10x easier now.
@Seb , here’s his implementation on the SqueezeExcite portion for reference:
def __init__(self, in_chs, reduce_chs=None, act_fn=F.relu, gate_fn=torch.sigmoid):
self.act_fn = act_fn
self.gate_fn = gate_fn
reduced_chs = reduce_chs or in_chs
self.conv_reduce = nn.Conv2d(in_chs, reduced_chs, 1, bias=True)
self.conv_expand = nn.Conv2d(reduced_chs, in_chs, 1, bias=True)
def forward(self, x):
# NOTE adaptiveavgpool can be used here, but seems to cause issues with NVIDIA AMP performance
x_se = x.view(x.size(0), x.size(1), -1).mean(-1).view(x.size(0), x.size(1), 1, 1)
x_se = self.conv_reduce(x_se)
x_se = self.act_fn(x_se)
x_se = self.conv_expand(x_se)
x = self.gate_fn(x_se) * x