Support for channels-first mode

[bryankim96]

When training on GPU, changing the dimension order can give a training speedup. It is unclear whether we are limited by I/O or by actual training speed, so the effect on final training performance may be limited.

See https://www.tensorflow.org/performance/performance_guide#data_formats for explanation.