qubvel / segmentation_models

Segmentation models with pretrained backbones. Keras and TensorFlow Keras.
MIT License
4.7k stars 1.03k forks source link

Training a simple U-Net without pretrained models #368

Open asmagen opened 4 years ago

asmagen commented 4 years ago

My application is segmentation of pathology images where the features might be very different from what has been used to derive the backbones and weights used here based on Imagenet. Is there a way to trained my own model from scratch based on the original U-Net architecture without the Imagenet component? Also, if I want to generate my own pretrained model based on publicly available pathology images, how can I configure it as a backbone for another segmentation task here as you do with Imagenet?

Thanks

JordanMakesMaps commented 4 years ago

"Is there a way to trained my own model from scratch based on the original U-Net architecture without the Imagenet component?"

Yes, just don't include pre-trained weights when initializing the model (e.g. include None instead of imagenet)

"... how can I configure it as a backbone for another segmentation task here as you do with Imagenet?"

@qubvel did a really great with the including different backbones (resnet-50, efficientnet-b4, etc) as encoders. If you were to train a resnet-50 for image classification, you could take those same weights and include them to the resnet-50 encoder within your Unet model, make sense?

Just to be clear (sorry if you already know this), but when you load ImageNet weights for say, a ResNet-50, those weights were made by training the same architecture on the dataset and you're just using them. But you can't use transfer the ImageNet weights for a ResNet50 to use with the DenseNet201: you have to use the ImageNet weights for a DenseNet201, and if they don't exist then you have to train it yourself.

So if you wanted to create weights that were learned from training a deep learning model on some public dataset for pathology, you could easily do that, and the weights for the encoder, the decoder, or both could be shared and used as pre-trained weights for when working with another pathology dataset. But those weights can only be used with the same architecture that you used previously.

asmagen commented 4 years ago

So to summarize, when I load a pre-trained network, it just means that I'm setting the backbone with specific weights rather than random weights. And the prediction output structure doesn't matter (say I trained on a 256x256 pixel tiles from a public pathology dataset with 10 classes to segment and I apply it on a 512x512 pixel tiles to predict 3 classes in my data, just as long I'm using the (None,None,None,3) setting for the input size). Is that correct?

asmagen commented 4 years ago

And I guess the other component of my question is just the backbone itself - I see here many (complex) backbones but I don't see the original U-Net backbone. The definition of the network here (model = sm.Unet('resnet34', classes=1, activation='sigmoid')) implies that the ResNet is a type of U-Net but I don't see how. How can I use the original structure of U-Net for semantic segmentation?

JordanMakesMaps commented 4 years ago

To your first reply: Yes. When it comes to the dimensions of the images, if you supply (None, None, None, 3) or however you have the order setup, then you can train and then test with different sized images. No issues should occur except that maybe, if you only train with one size/scale and then test on a completely different size/scale, then you might have incorrect predictions, So make sure that your training data is a good representation of what you are likely to see in the wild and/or make sure that you pre-process images from the wild. Regarding the classes, if you initially train a model with say 10 classes, it will only know those 10 classes. If you have a new dataset with something other than those 10 classes (maybe an 11th class), then it will try to make predictions on the 11th class but obviously get them wrong. To incorporate an 11th, 12th, etc, class, you'll just use the weights as pre-trained and train on top of them with all of the classes of interest.

Second reply: yes, this repo only has known backbones as options as encoders (resnet-34, densenet-201, etc). If you want to have what people commonly refer to as a "Vanilla" model, then you'll need to create it yourself or find another repo that has one (like this one). But, I'll say that this repo that uses the complex backbones tend to work better.

asmagen commented 4 years ago

Thanks. Somehow the model definition doesn't accept the None parameter, only imagenet.

# define network parameters
n_classes = 1 if len(SEGMENTATION_CLASSES) == 1 else (len(SEGMENTATION_CLASSES) + 1)  # case for binary and multiclass segmentation
activation = 'sigmoid' if n_classes == 1 else 'softmax'

#create model
model = sm.Unet(BACKBONE, classes=n_classes, activation=activation, encoder_weights='None',input_shape = (None, None, 3))
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-31-a0481d301d20> in <module>()
      4 
      5 #create model
----> 6 model = sm.Unet(BACKBONE, classes=n_classes, activation=activation, encoder_weights='None',input_shape = (None, None, 3))

5 frames
/usr/local/lib/python3.6/dist-packages/keras_applications/densenet.py in DenseNet(blocks, include_top, weights, input_tensor, input_shape, pooling, classes, **kwargs)
    178 
    179     if not (weights in {'imagenet', None} or os.path.exists(weights)):
--> 180         raise ValueError('The `weights` argument should be either '
    181                          '`None` (random initialization), `imagenet` '
    182                          '(pre-training on ImageNet), '

ValueError: The `weights` argument should be either `None` (random initialization), `imagenet` (pre-training on ImageNet), or the path to the weights file to be loaded.

Screen Shot 2020-07-04 at 8 32 17 AM

JordanMakesMaps commented 4 years ago

In python None is a keyword, so you don't need to put quotes around it

None != 'None'