Improve the training procedure

[vincekurtz]

The current training loop is fairly efficient, but lacks some features. Eventually we will want:

• [ ] A proper validation set (should be generated independently, since data points are correlated)
• [ ] Logging (to tensorboard or similar)
• [ ] Input normalization