Support importing train.py

[amosonn]

Motivation

To allow running the training tool from within a python session (and not just as a commandline script)

Modification

Make the "main" function more accessible (supporting arguments)

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}