OOM error when training on my own dataset about segcaps HOT 4 CLOSED

Hi CreepZzy,

The model did not test on 2D grayscale images so you may need to modify the code.
For memory issue, my suggestion is to reduce the batch size or to restrict tensorflow on GPU usage.

Good Luck to you.

With kind regards,
Cheng-Lin Li

from segcaps.

Hello @CreepZzy,

Unfortunately at 512 x 512 input sizes, with the current SegCaps structure, this takes up almost the entire memory of a 12GB GPU. My suggestions would be to either change the code to accept something like 256 x 256 inputs or to change the network of SegCaps to be lighter (less capsule types per layer or less layers). Hope this helps. For the test output, just change the plotting within test.py to whatever colors you want.

from segcaps.

@lalonderodney Thank you for your helpful reply. I'll have try on these kinds of modification. I don't know that CapsNet can be such GPU hunger.

from segcaps.

@CreepZzy,

In order to do the dynamic routing of child capsules to parent capsules, you must store the intermediate representations to be routed. This takes memory. It's similar to how DenseNets, although having far fewer parameters take more memory. Somewhat recently many implementations have come out for memory efficient implementations of DenseNets. Hopefully something similar will be eventually developed by intelligent people for capsules :)

from segcaps.