in part_generator.py [0:0]
def __init__(self, name, results_dir, models_dir, n_part, image_size, network_capacity, batch_size = 4, mixed_prob = 0.9,
gradient_accumulate_every=1, lr_D = 2e-4, lr_G = 2e-4, num_workers = None, save_every = 1000, trunc_psi = 0.6, sparsity_penalty=0.):
self.GAN = None
self.name = name
self.results_dir = Path(results_dir)
self.models_dir = Path(models_dir)
self.config_path = self.models_dir / name / '.config.json'
assert log2(image_size).is_integer(), 'image size must be a power of 2 (64, 128, 256, 512, 1024)'
self.n_part = n_part
self.image_size = image_size
self.network_capacity = network_capacity
self.lr_D = lr_D
self.lr_G = lr_G
self.batch_size = batch_size
self.num_workers = num_workers
self.mixed_prob = mixed_prob
self.sparsity_penalty = sparsity_penalty
self.save_every = save_every
self.steps = 0
self.trunc_psi = trunc_psi
self.gradient_accumulate_every = gradient_accumulate_every
self.d_loss = 0
self.g_loss = 0
self.last_gp_loss = 0
self.pl_loss = 0
self.sparsity_loss = 0
self.pl_mean = torch.empty(1).cuda()
self.pl_ema_decay = 0.99
self.loader_D = None
self.loader_G = None
self.av = None
if 'bird' in self.name:
self.part_to_id = {'initial': 0, 'eye': 1, 'head': 4, 'body': 3, 'beak': 2, 'legs': 5, 'wings': 8, 'mouth': 6, 'tail': 7}
COLORS = COLORS_BIRD
elif 'generic' in self.name or 'fin' in self.name or 'horn' in self.name:
self.part_to_id = {'initial': 0, 'eye': 1, 'arms': 2, 'beak': 3, 'mouth': 4, 'body': 5, 'ears': 6, 'feet': 7, 'fin': 8,
'hair': 9, 'hands': 10, 'head': 11, 'horns': 12, 'legs': 13, 'nose': 14, 'paws': 15, 'tail': 16, 'wings':17}
COLORS = COLORS_GENERIC
self.color = 1-torch.cuda.FloatTensor([0, 0, 0]).view(1, -1, 1, 1)
self.default_color = 1-torch.cuda.FloatTensor([0, 0, 0]).view(1, -1, 1, 1)
for key in COLORS:
if key in self.name:
self.color = COLORS[key]
break
for partname in self.part_to_id.keys():
if partname in self.name:
self.partid = self.part_to_id[partname]
self.partname = partname