Artifact Model Training Optimizations

[fatsmcgee]

All performance numbers cited below are on an 8 CPU machine (n1-standard-8) with a Tesla T4 GPU.

This PR introduces the following:

• For data iteration / inference outside of model learning, which includes both model evaluation and the construction of ArtifactDataset, allow for distinct inference batch sizes (larger = much faster) and use torch inference mode. When inference batch size is large (default is set to 8192), this makes inference much faster (E.g ArtifactDataset construction goes from 6.5 to 3.5 minutes).
• Resolve performance problems with num_workers>0 which were most likely caused by issues with serialization of ArtifactBatch.original_data in DataLoader works. Instead of storing original_data, only store the pieces of that data which are accessed.
  • Before with GPUs, num_workers=4 was much slower than num_workers=0. Now it is a bit faster (8 mins 20 second epoch versus 8 mins 34).
  • This also resolves a runtime crash that occurs up in Linux environments where the open file limit is low (I observed this on a Google Cloud VM). This is probably because PyTorch opens shared memory "files" to communicate data between workers, and original_data had lots of nested Python (not PyTorch/Numpy) data which got serialized independenty.

[davidbenjamin]

This will also need to go in the WDL scripts, but I'll handle that in a follow-up PR.