Improve train_generator script.

[macfergus]

1. Checkpoint weights after each epoch. Good for long-running jobs that converge early.
2. Change path to the model_zoo so it does not depend on the cwd.
3. Make the sample size and nb_epochs command-line args.

[maxpumperla]

Thanks, that should prove very useful!