Retraining a brainy model (unet) on kwyk dataset using nobrainer API

[hvgazula]

Based on a user request #271

[hvgazula]

• setup kwyk data
• unet model (nobrainer library model)
• unet model (standard tf model)
• environment setup
• training on dgx-100 (volta) -- fails
• training on a100 - success with batch size = 1

[satra]

@hvgazula - the request was about retraining the brainy model not the kwyk model (you may want to edit the title). the kwyk model uses the bayesian neural network. i would love to see that trained as well, but that's a different task.

[hvgazula]

Setting block_shape = None runs successfully on the test data (whose shape is (256, 256, 256). However, the process gets killed while running the same on the kwyk tfrecords despite setting the batch size to 1/2 (depending on the number of GPUs). This may have to do with how huge the tfrecord file is. To test, create a new tfrecord file with fewer shards and see if that works.

[hvgazula]

Note: The successful run was on 2 A100s.