During my optimization loop I have the need to visualize a lot of my training output using matplotlib and other tools. Hence, I need to copy data from GPU to CPU and subsequently use matplotlib to plot (and save) data which takes quite a while.
This became so exhaustive that a large amount of time is now spent on plotting these intermediate results while the GPU utilization drops down to 0. I’ve been looking into flags like non_blocking, pin_memory etc., but they seem to focus mostly on CPU → GPU memory transfers.
Is there any way that plotting and training can be done asynchronously ? Ideally, I’d like to continue training on the GPU while the CPU is used to do “minor tasks” like plotting, saving on disk etc.
Thanks