I try to do as much of the transformations as possible in tensorflow using the Datasets API so that I don't have to write them in another language/system in production.
I had tried to do the same. However given the CPU/GPU imbalance on AWS GPU instances I have resorted to building a fully "rendered" training set and doing all the transforms in spark.
See: https://github.com/tensorflow/tensorflow/issues/13610