We support use different dataset(imagenet, cifar10 and mnist) to train and test the model. But the code is hard to maintain. We should do some code refactor.
The target is unify the dataload, transform function of different dataset in local or spark cluster mode.