Uses EfficientNetB0 as segmentation model encoder backbone

[daniel-j-h]

For #172 (see for context) - this implements the EfficientNetB0 model as an encoder for our encoder-decoder architecture.

I'm currently training my EfficientNet model family (no h-swish, no squeeze-and-excitation) in

https://github.com/daniel-j-h/efficientnet

on ImageNet and want to see how they behave as backbone encoder in robosat. I'm inlining the EfficientNet implementation here without the h-swish, scSE, or implant code.

The encoder blocks are tiny compared to the previous ResNet50 blocks. If the EfficientNetB0 features are strong enough we might want to spend some of the resources we gained in the decoder blocks, e.g. res-blocks for learned upsampling or PixelShuffle+ICNR init for learned upsampling, scSE blocks, or simply more features.

Needs thorough evaluation before merging; mainly opening this for visibility.

cc @ocourtin

[daniel-j-h]

Here's the core of this changeset

robosat/robosat/unet.py

Lines 122 to 157 in 4a3c123

	# 1, 3, 512, 512
	enc0 = self.net.conv1(x)
	enc0 = self.net.bn1(enc0)
	enc0 = self.net.relu(enc0)
	# 1, 32, 256, 256
	enc0 = self.net.layer1(enc0)
	# 1, 16, 256, 256
	enc1 = self.net.layer2(enc0)
	# 1, 24, 128, 128
	enc2 = self.net.layer3(enc1)
	# 1, 40, 64, 64
	enc3 = self.net.layer4(enc2)
	# 1, 80, 32, 32
	enc3 = self.net.layer5(enc3)
	# 1, 112, 32, 32
	enc4 = self.net.layer6(enc3)
	# 1, 192, 16, 16
	enc4 = self.net.layer7(enc4)
	# 1, 320, 16, 16
	enc4 = self.net.features(enc4)
	# 1, 1280, 16, 16
	center = self.center(nn.functional.max_pool2d(enc4, kernel_size=2, stride=2))
	dec0 = self.dec0(torch.cat([enc4, center], dim=1))
	dec1 = self.dec1(torch.cat([enc3, dec0], dim=1))
	dec2 = self.dec2(torch.cat([enc2, dec1], dim=1))
	dec3 = self.dec3(torch.cat([enc1, dec2], dim=1))
	dec4 = self.dec4(dec3)
	dec5 = self.dec5(dec4)
	return self.final(dec5)

robosat/robosat/efficientnet.py

Lines 91 to 103 in 4a3c123

	self.conv1 = nn.Conv2d(3, scaled_width(32), kernel_size=3, stride=2, padding=1, bias=False)
	self.bn1 = nn.BatchNorm2d(scaled_width(32))
	self.relu = nn.ReLU6(inplace=True)
	self.layer1 = self._make_layer(n=scaled_depth(1), expansion=1, cin=scaled_width(32), cout=scaled_width(16), kernel_size=3, stride=1)
	self.layer2 = self._make_layer(n=scaled_depth(2), expansion=6, cin=scaled_width(16), cout=scaled_width(24), kernel_size=3, stride=2)
	self.layer3 = self._make_layer(n=scaled_depth(2), expansion=6, cin=scaled_width(24), cout=scaled_width(40), kernel_size=5, stride=2)
	self.layer4 = self._make_layer(n=scaled_depth(3), expansion=6, cin=scaled_width(40), cout=scaled_width(80), kernel_size=3, stride=2)
	self.layer5 = self._make_layer(n=scaled_depth(3), expansion=6, cin=scaled_width(80), cout=scaled_width(112), kernel_size=5, stride=1)
	self.layer6 = self._make_layer(n=scaled_depth(4), expansion=6, cin=scaled_width(112), cout=scaled_width(192), kernel_size=5, stride=2)
	self.layer7 = self._make_layer(n=scaled_depth(1), expansion=6, cin=scaled_width(192), cout=scaled_width(320), kernel_size=3, stride=1)
	self.features = nn.Conv2d(scaled_width(320), scaled_width(1280), kernel_size=1, bias=False)