DynamicUnet error

fastai fastai dev
1

Hi,
I would like to first fine-tune a resnet34 classifier on a task and then use it as the encoder of a U-Net model for segmentation purposes.
I proceeded as follows:

  • Fine-tune resnet34
    learner= cnn_learder(data, models.resnet34, ...)
    learner.fit(epochs=1)

  • Remove the head
    model = learner.model
    removed = list(model.children())[:-1]
    model= torch.nn.Sequential(*removed)

  • Create U-Net
    DynamicUnet(model, n_classes=2)

It returns the following error:

image
image1532×596 66.7 KB

I suspect that DynamicUnet does not like the cut

2

I’ve found the solution quicker than expected
After cutting the head the model has a sequential block which encapsulates the targeted model:
model = learner.model
model = DynamicUnet(model[0].to('cpu'), n_classes=1)
is working fine. I also had to move the model to CPU to build U-Net

1 Like