refactoring train-test code
Merge pull request #1 from riktor/master

refactoring train-test code