search

sjeknic / CellTK

Toolkit for analysis of live-cell microscopy data
MIT License
3 stars 0 forks source link

`UNet` is broken compared to previous versions #76

Closed sjeknic closed 2 years ago

sjeknic commented 2 years ago

This could end up being a huge pain. The weights that I have that I used with CellUNet no longer seem to produce that same results. Nuclear segmentation is different (I think), but useable. Cytoplasm is just a shit show. See below, old one is first: old new

Clearly something is wrong. It seems likely that the model was built correctly, because otherwise I don't think the weights would have loaded properly. So could this be a backwards compatibility issue in the packages. That would be awful, because I would need to recreate the training data and also add prediction (#59). That isn't super simple though, as I need to add augmentation batch picking, making the training data generator, making a separate test data generator. Overall, just a ton of painful work...

sjeknic commented 2 years ago

Isn't it fun when I memorialize my mistakes in perpetuity.

This wasn't caused by an issue in UNet, but rather an issue in picking which files get passed where. I believe that it is fixed in #75

sjeknic commented 2 years ago

Okay, need to reopen this, haha. See below. Seems to be an issue related to the properly handling intermediate values. I think this is likely a dtype issue. Let's see. Again, old is first. old new

sjeknic commented 2 years ago

I think I have this fixed now in #75. There were a couple issues.

First, I had, for some reason, included a softmax activation layer after the last convolutional layer. While this isn't necessarily a terrible design choice, it does not match the UNet paper, nor the implementation in CellTK. Instead, there actually isn't any activation layer after the last convolutional layer.

This brings us to the second issue. The final convolutional layer in this repo didn't have any options passed to it besides the number of filters (3 for number of regions) and kernel size (1). This means that it will default to activation 'relu' and, more important for this issue, padding mode of 'border'. This is in contrast to CellTK, which explicitly sets the activation and padding (well border_mode because it's old). The padding should be 'same', not 'border'.

Finally, there was also a small mistake in the function that calculated the padded dimensions. The model must be padded to be divisible by 8 (determined by model depth) in each axis. However, if the axis was already divisible by 8, it would still get padded (by 4 on each side). So the model would run just fine, but especially close to the left and right edges, the layer would be getting extra values passed to it. This is also now fixed.

As far as I can tell, the outputs are now identical with CellTK both before and after normalization.