How to use unet_learner with only 1 channel?

fastai
1

Hi all, a continuation of my previous attempt (still a newbie now moving on to lesson4). Now trying to get the unet_learner to run with 1 channel input.
For reference, both my masks and my images are simple 1channel .tif files.
image

Code I am running, arch= various 
import fastbook
fastbook.setup_book()
from fastbook import *
path =Path("D:/pytorch/data/2D_Zebrafish/fastai")
codes = ['Zebrafish']

dls = SegmentationDataLoaders.from_label_func(
    path, bs=12, fnames = get_image_files(path/"images"),
    label_func = lambda o: path/'labels'/f'{o.stem}_annotationLabels.tif',
    codes =codes, num_workers=0
)

def custom_resnet18(*args, **kwargs):
    model = resnet18(*args, **kwargs)
    model.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
    return model

#####ERROR WHY NIN=3####
learn = unet_learner(dls, arch = ???, n_in=1, n_out=1, loss_func=MSELossFlat(), lr=0.0001, normalize=False)
learn.fine_tune(1)

I am getting an error similar to this post, but following the instructions there for a custom model does not change the result. I also am unable to load another pretrained unet using:

model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch', 'unet', in_channels=1, out_channels=1, init_features=8, pretrained=False)

I can get that model to work for my data using standard PyTorch. I would really like to learn to use FastAI, though.

Plugging resnet18 into the arch gives me:

RuntimeError: Given groups=1, weight of size [64, 1, 7, 7], expected input[12, 3, 512, 512] to have 1 channels, but got 3 channels instead

While plugging in the custom_resnet18 defined above and in the linked post gives:

AssertionError: Unexpected number of input channels, found 1 while expecting 3

That is the same as the original error in the above post, so that thread did not resolve the issue for me.

Full error 
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
Input In [39], in <cell line: 19>()
     16     return model
     18 #####ERROR WHY NIN=3####
---> 19 learn = unet_learner(dls, custom_resnet18, n_in=1, n_out=1, loss_func=MSELossFlat(), lr=0.0001, normalize=False)
     20 learn.fine_tune(1)

File D:\Anaconda\envs\fastai2\lib\site-packages\fastai\vision\learner.py:245, in unet_learner(dls, arch, normalize, n_out, pretrained, config, loss_func, opt_func, lr, splitter, cbs, metrics, path, model_dir, wd, wd_bn_bias, train_bn, moms, **kwargs)
    243 img_size = dls.one_batch()[0].shape[-2:]
    244 assert img_size, "image size could not be inferred from data"
--> 245 model = create_unet_model(arch, n_out, img_size, pretrained=pretrained, **kwargs)
    247 splitter=ifnone(splitter, meta['split'])
    248 learn = Learner(dls=dls, model=model, loss_func=loss_func, opt_func=opt_func, lr=lr, splitter=splitter, cbs=cbs,
    249                metrics=metrics, path=path, model_dir=model_dir, wd=wd, wd_bn_bias=wd_bn_bias, train_bn=train_bn,
    250                moms=moms)

File D:\Anaconda\envs\fastai2\lib\site-packages\fastai\vision\learner.py:220, in create_unet_model(arch, n_out, img_size, pretrained, cut, n_in, **kwargs)
    218 "Create custom unet architecture"
    219 meta = model_meta.get(arch, _default_meta)
--> 220 body = create_body(arch, n_in, pretrained, ifnone(cut, meta['cut']))
    221 model = models.unet.DynamicUnet(body, n_out, img_size, **kwargs)
    222 return model

File D:\Anaconda\envs\fastai2\lib\site-packages\fastai\vision\learner.py:77, in create_body(arch, n_in, pretrained, cut)
     75 "Cut off the body of a typically pretrained `arch` as determined by `cut`"
     76 model = arch(pretrained=pretrained)
---> 77 _update_first_layer(model, n_in, pretrained)
     78 if cut is None:
     79     ll = list(enumerate(model.children()))

File D:\Anaconda\envs\fastai2\lib\site-packages\fastai\vision\learner.py:57, in _update_first_layer(model, n_in, pretrained)
     55 first_layer, parent, name = _get_first_layer(model)
     56 assert isinstance(first_layer, nn.Conv2d), f'Change of input channels only supported with Conv2d, found {first_layer.__class__.__name__}'
---> 57 assert getattr(first_layer, 'in_channels') == 3, f'Unexpected number of input channels, found {getattr(first_layer, "in_channels")} while expecting 3'
     58 params = {attr:getattr(first_layer, attr) for attr in 'out_channels kernel_size stride padding dilation groups padding_mode'.split()}
     59 params['bias'] = getattr(first_layer, 'bias') is not None

AssertionError: Unexpected number of input channels, found 1 while expecting 3

I thought from the post above that the unet_learner should be able to handle single channel inputs by adding the weights. How do I make that work? I am fine with either using the pretrained resnet or figuring out how to load in the pretrained model.

This seems very related, but I am still trying to work out how to actually use it: UNet: Size error for a custom dataset - #29 by cordmaur

2

Hello,

Setting arch to any model that ships with fastai, e.g., resnet18, works fine; the error below is indicating the model expects a single input channel but is being provided with 3. That is, your data is being processed incorrectly and contains 3 channels instead of 1 channel, an issue that can be traced to fastai’s SegmentationDataLoaders.from_label_func because it converts gray-scale pictures to RGB ones. To get the intended behaviour, you can implement a custom DataBlock that ensures the input is being loaded properly.

Using custom_resnet18 and n_in = 1 together leads to the following error since fastai assumes the model’s input channels is 3, and the user would like to modify it to something else. Removing n_in = 1 would solve it.

Finally, for using the PyTorch UNet you have linked with fastai, you would wrap it in a Learner alongside your data to acquire fastai functionalities.

Please let me know if you have any other questions.

2 Likes
3

Thanks, I think that is what I need to look into doing. I did end up getting things working either by going with pure PyTorch or simply allowing the conversion to greyscale RGB, but since I would like to scale this, I would prefer to understand how to choose the correct number of channels (which may someday be 2 or 4, 3D, etc.) and load in pretrained weights for a given number of independent channels. Much more to learn.

Yep, think I had that in the code above. Thanks.

1 Like