faceswap-GAN icon indicating copy to clipboard operation
faceswap-GAN copied to clipboard

Use prefetch_generator to increase GPU utilization

Open slundell opened this issue 6 years ago • 1 comments

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,:,:,:]       

slundell avatar Mar 03 '18 18:03 slundell

Thanks for the tips!

shaoanlu avatar Mar 07 '18 06:03 shaoanlu