change pre-trained Model for fine-tuning does not work

ManuelNiemeyer — Tue, 12 Sep 2023 09:20:02 GMT

Hello,
I would like to use the following pre-trained model for a fine-tuning.

https://www.esri.com/arcgis-blog/products/arcgis-pro/imagery/deep-learning-with-arcgis-pro-tips-tricks-part-2/

However, the model was trained on RGB aerial images and I have RGBI aerial images.
Trying to use this model as a basis causes the following error:

RuntimeError: Error(s) in loading state_dict for DynamicUnet:
	size mismatch for layers.0.0.weight: copying a param with shape torch.Size([64, 3, 7, 7]) from checkpoint, the shape in current model is torch.Size([64, 4, 7, 7]).
	size mismatch for layers.10.layers.0.0.weight: copying a param with shape torch.Size([49, 99, 3, 3]) from checkpoint, the shape in current model is torch.Size([50, 100, 3, 3]).
	size mismatch for layers.10.layers.0.0.bias: copying a param with shape torch.Size([49]) from checkpoint, the shape in current model is torch.Size([50]).
	size mismatch for layers.10.layers.1.0.weight: copying a param with shape torch.Size([99, 49, 3, 3]) from checkpoint, the shape in current model is torch.Size([100, 50, 3, 3]).
	size mismatch for layers.10.layers.1.0.bias: copying a param with shape torch.Size([99]) from checkpoint, the shape in current model is torch.Size([100]).
	size mismatch for layers.11.0.weight: copying a param with shape torch.Size([9, 99, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 100, 1, 1]).
	size mismatch for layers.11.0.bias: copying a param with shape torch.Size([9]) from checkpoint, the shape in current model is torch.Size([3]).

The only way to solve this problem was to modify the first layer as follows:

class CustomUNet(nn.Module):
def __init__(self):
super(CustomUNet, self).__init__()
# Change first Layer from three to four input channels
self.conv1 = nn.Conv2d(4, 64, kernel_size=7, stride=2, padding=3, bias=False)
# add the other layers from the pre-trained model
self.pretrained = unet_pretrained

def forward(self, x):
x = self.conv1(x)
x = self.pretrained(x)
return x

custom_UNet = CustomUNet()

However, I have not been able to convert the model back to a UnetClassifier afterwards.

Can anyone help me?

topic change pre-trained Model for fine-tuning does not work in Developers Questions

change pre-trained Model for fine-tuning does not work