faceswap-GAN
faceswap-GAN copied to clipboard
Use prefetch_generator to increase GPU utilization
Not really an issue, more of a tip.
Generating a batch takes around 20% of the iteration time on my setup. Since training is CPU/GPU bound and batch assembly is IO bound, performance can be gained by doing this asynchronously. The easiest way I found is to use prefetch_generator https://github.com/justheuristic/prefetch_generator This setup does not support .send() which is used only to set the preview batch size. I can do with preview batch size being the same as the training batch size. By removing the .send() and making minibatch a prefetch_generator my GPU is utilized at 100% almost all the time. My iteration time went from around 2300ms to 1850ms.
@background(36)
def minibatch(data, size):
length = len(data)
epoch = i = 0
shuffle(data)
while True:
if i+size > length:
shuffle(data)
i = 0
epoch+=1
rtn = np.float32([read_image(data[j]) for j in range(i,i+size)])
i+=size
yield epoch, rtn[:,0,:,:,:], rtn[:,1,:,:,:]
Thanks for the tips!