bfortuner / pytorch_tiramisu

FC-DenseNet in PyTorch for Semantic Segmentation
MIT License
306 stars 65 forks source link

Take portion of tiramisu #5

Closed varghesealex90 closed 6 years ago

varghesealex90 commented 6 years ago

Hi;

After training a tiramisu, I would like to just use all the layers till the bottleneck layer for feature extraction. Is there a good and recommended way to do it. ?

bfortuner commented 6 years ago

Yep! You can "cut" the model object after training like this to extract the feature maps you're looking for:

https://github.com/bfortuner/pytorch-kaggle-starter/blob/master/models/builder.py#L40

the 'cut' param is the index of the layer