Spaces:
Runtime error
Runtime error
File size: 2,349 Bytes
c2f174c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import random
import torch
class LatentCodesPool:
"""This class implements latent codes buffer that stores previously generated w latent codes.
This buffer enables us to update discriminators using a history of generated w's
rather than the ones produced by the latest encoder.
"""
def __init__(self, pool_size):
"""Initialize the ImagePool class
Parameters:
pool_size (int) -- the size of image buffer, if pool_size=0, no buffer will be created
"""
self.pool_size = pool_size
if self.pool_size > 0: # create an empty pool
self.num_ws = 0
self.ws = []
def query(self, ws):
"""Return w's from the pool.
Parameters:
ws: the latest generated w's from the generator
Returns w's from the buffer.
By 50/100, the buffer will return input w's.
By 50/100, the buffer will return w's previously stored in the buffer,
and insert the current w's to the buffer.
"""
if self.pool_size == 0: # if the buffer size is 0, do nothing
return ws
return_ws = []
for w in ws: # ws.shape: (batch, 512) or (batch, n_latent, 512)
# w = torch.unsqueeze(image.data, 0)
if w.ndim == 2:
i = random.randint(0, len(w) - 1) # apply a random latent index as a candidate
w = w[i]
self.handle_w(w, return_ws)
return_ws = torch.stack(return_ws, 0) # collect all the images and return
return return_ws
def handle_w(self, w, return_ws):
if self.num_ws < self.pool_size: # if the buffer is not full; keep inserting current codes to the buffer
self.num_ws = self.num_ws + 1
self.ws.append(w)
return_ws.append(w)
else:
p = random.uniform(0, 1)
if p > 0.5: # by 50% chance, the buffer will return a previously stored latent code, and insert the current code into the buffer
random_id = random.randint(0, self.pool_size - 1) # randint is inclusive
tmp = self.ws[random_id].clone()
self.ws[random_id] = w
return_ws.append(tmp)
else: # by another 50% chance, the buffer will return the current image
return_ws.append(w)
|