|
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: |
|
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: |
|
return ws |
|
return_ws = [] |
|
for w in ws: |
|
|
|
if w.ndim == 2: |
|
i = random.randint(0, len(w) - 1) |
|
w = w[i] |
|
self.handle_w(w, return_ws) |
|
return_ws = torch.stack(return_ws, 0) |
|
return return_ws |
|
|
|
def handle_w(self, w, return_ws): |
|
if self.num_ws < self.pool_size: |
|
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: |
|
random_id = random.randint(0, self.pool_size - 1) |
|
tmp = self.ws[random_id].clone() |
|
self.ws[random_id] = w |
|
return_ws.append(tmp) |
|
else: |
|
return_ws.append(w) |
|
|