search

rgmyr / corebreakout

Segmentation and depth-alignment of geological core sample image columns via Mask-RCNN
Other
24 stars 15 forks source link

Pre-trained model and weights in assets.zip have different shape? #18

Closed metazool closed 4 years ago

metazool commented 4 years ago

Hello! Thank you for your work on corebreakout. It's wonderful to see BGS open licenced data being used in ML projects.

I've been trying to use the CoreSegmenter interface on some more of our imagery, however when i try loading it with the pre-trained model and weights available in assets.zip I see this error about the weights and the model having different shapes:

ValueError: Layer #389 (named "mrcnn_bbox_fc"), weight <tf.Variable 'mrcnn_bbox_fc/kernel:0' shape=(1024, 12) dtype=float32_ref> has shape (1024, 12), but the saved weight has shape (1024, 324).

Here it is being thrown in a test: https://github.com/metazool/corebreakout-fastapi/runs/997750099?check_suite_focus=true#step:5:473

This is what I'm providing for model_dir and weights_path, maybe I have missed something obvious? https://github.com/metazool/corebreakout-fastapi/blob/main/coreapi/config.py

The idea here is a proof of concept for another project where we are looking in more detail at the extent of core fragmentation and degradation in the older boxes, I hope we will supply another set of labels for the training data you are already using, and take advantage of corebreakout for some of the preprocessing, so I really appreciate being able to build on this work

mobiuscreek commented 4 years ago

Hello and thank you for a very nice and useful package. The error is caused because the default config has two + 1 classes and not 81 classes that the COCO pre-trained network has and that causes a weight dimension mismatch. To solve this, some layers have to be excluded.

Relevant to the issue, in the test_inference.ipynb notebook provided, the CoreSegmenter interface is used to make inference. That's because it defaults to inference mode as seen here: https://github.com/rgmyr/corebreakout/blob/d2a6ea951805aa235eed0f4897132b5459d97d15/corebreakout/segmenter.py#L68-L69 it then goes to load the weights and that causes the error.

However, the MaskRCNN interface is used for training in the train_mrcnn_model.ipynb notebook. Would it make sense to subclass MaskRCNN instead so the same interface can be used for both inference and training?

If the CoreSegmenter mode was configurable, this behaviour and the error could have been prevented.

zanejobe commented 4 years ago

Hello @metazool and @mobiuscreek ,

Thanks for testing it out, and helping us improve! @rgmyr is looking into this, and will update you soon.

We'd also appreciate a shout-out on the BGS CoreScore ReadMe if you have learned from this repo - probably best to use the JOSS paper, which has more explanation https://doi.org/10.21105/joss.01969

Thanks! Zane

metazool commented 4 years ago

Thanks @zanejobe - the CoreScore project is an MSc project from a couple of years ago which we've recently taken over in an attempt to productionise for QA of material in the BGS core store. We haven't been able to make use of corebreakout directly yet but are definitely hoping there will be ways to make the two projects work together. If citations and shoutouts would help I am more than happy to offer them. It was the JOSS paper that brought this project onto our radar, good work on that :)

rgmyr commented 4 years ago

Hi @metazool, thanks for the interest, and apologies for the wait. @mobiuscreek is correct in that the CoreSegmenter is setup for inference only, not for training.

The idea was that you would train a model using scripts/train_mrcnn_model.py or notebooks/train_mrcnn_model.ipynb, and then you should be able to instantiate a CoreSegmenter with those trained weights. You should also be able to load assets/models/mask_rcnn_cb_default.h5 into a CoreSegmenter without doing any training.

I agree that this isn't the cleanest interface, although I'm a little biased against subclassing mrcnn.MaskRCNN since that would pull a lot of extra methods into the class that I don't anticipate being relevant to the user. Basically composition over inheritance.

I think the way I'd prefer to achieve this flexibility would be to add a training interface to CoreSegmenter, although I'm concerned that mrcnn is getting pretty stale (1.5 years since last commit, 2.5 since last release), and I'm reluctant to invest much effort in extending it's use here (it's always been a little buggy to begin with).

I had considered converting to Detectron2 since I've mostly switched to pytorch and it's a very popular and active project. For now I don't quite have the time to commit to that large of a refactor, but I'm open to feedback on what would be useful for you and what you might like to see improved in the short term.

metazool commented 4 years ago

Detectron2 does look interesting, thank you. It would be a nicer path than updating to Tensorflow 2 especially if mask-rcnn is not being actively maintained. In the short term we're committed to working with the Unet implementation in fastai with the project we inherited. I really hope we can come back to the conversation and converge in the future.

Thank you for looking into this though! We can't load mask_rcnn_cb_default.h5 into the Segmenter with the provided weights due to the issue described above. Retraining from scratch as you suggest is probably the easiest path.