Re-use U-Net for Transfer Learning (how to "chop off" layers)?

Hi all,

I’ve trained a U-Net model that does segmentation, but now I would like to play with the first half (encoder) of the network. For example, to use the middle layer (output of the first half) as an embedding, or even using the first half as initial weights for transfer learning on other tasks like image classification. How do I do that?

I searched around, and it seemed like in 2018, you could do stuff like this:

custom_head = nn.Sequential(AdaptiveConcatPool2d(), Flatten())
new_model = nn.Sequential(*list(children(old_model))[:-20], custom_head)

Which would chop off the last 20 layers, and add a new custom_head, which in this case just pool and flattens out. But in the current version the [:20] doesn’t seem to work. It looks like list(children(model)) has a length of 1?

Ok, it looks like in the 2019 version this seems to work (note the [] references the model object and not the list):

custom_head = nn.Sequential(AdaptiveConcatPool2d(), Flatten())
new_model = nn.Sequential(*list(children(old_model[:-20])), custom_head)

However, I can’t use the new_model object to create a new cnn, using cnn_learner. I can’t seem to find any material on how to re-use a segment of fastai trained networks, for transfer learning?

If anyone has useful resources, notebooks of how to manipulate trained fastai learner model objects, that would be much appreciated. Thanks!

1 Like