pytorch backend major update

[farakiko]

Checks and updates to the pytorch code,

• Fix the jet reconstruction and inference code
• Run multiple validation runs within a single epoch to avoid waiting for a full epoch
• Avoid putting the loss values on cpu() after every training step
• Make --num-workers configurable from the args
• Other formatting

The PR is tested through

1. Training on 100Top and 100QCD events for 300 epochs on 4 NVIDIA GeForce GTX 1080 Ti Gpus using the following command (eta ~90min):
   python mlpf/pyg_pipeline.py --dataset cms --config parameters/pyg-cms-small.yaml --gpus "0,1,2,3" --prefix /pfvol/experiments/MLPF_cms_small_PR_ --conv-type gravnet --num-epochs 300 --train --ntrain 100 --nvalid 100 --gpu-batch-multiplier 5

2. Testing on 500 high-pT QCD events 4 NVIDIA GeForce GTX 1080 Ti Gpus using the following command (eta ~18min):
   python mlpf/pyg_pipeline.py --dataset cms --config parameters/pyg-cms-test-qcdhighpt.yaml --gpus "0,1,2,3" --load "" --test --ntest 500 --make-plots --gpu-batch-multiplier 5

Results are shown below,
mlpf_loss_Total.pdf
jet_res.pdf
met_res.pdf

A few notes for the next update regarding multgpu training

• Must fix num-workers>0 for gpus>1 (atm runs into error)
• Broadcast stale_epochs to all gpus to stop the training in an optimal way (the code hangs at the moment when stale_epochs>patience)
• Careful look at dist.barrier() invokes during training to avoid bottlenecks

[jpata]

@erwulff could you take a look at the PR also?