Sure thing, here is the function that is calling:
class DiscGradReg(nn.Module):
"""Discriminator gradient regularization for inversion"""
def __init__(self, gamma: float = 10.):
"""
Args:
net: String representing the network to use for LPIPS loss (default: alex)
"""
super().__init__()
self.gamma = gamma
def forward(self, inpt, output):
# Disc represents the discriminator loss
return (self.gamma / 2.0) * (torch.norm(grad(output.mean(), inpt)[0]) ** 2)
and this is the main training loop
def search(**kwargs):
device = torch.device('cuda')
opts = dnnlib.EasyDict(kwargs) # Command line arguments.
# Determine the number of classes
num_classes = 0
for folders in os.listdir(opts['data']):
num_classes += 1
# Define transforms
tsfms = transforms.Compose([
transforms.Resize(256, interpolation=1),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
# Calculate number of steps to accumulate gradient
accumulate_steps = (opts['batch'] // opts['batch_gpu']) if opts.get('batch_gpu') else 1
dataset = ImageFolder(opts['data'], transform=tsfms, target_transform=lambda x: F.one_hot(torch.tensor(x, dtype=torch.int64), num_classes).float())
dataloader = DataLoader(
dataset,
shuffle=True,
batch_size=opts['batch_gpu'] if opts.get('batch_gpu') else opts['batch'],
drop_last=True,
num_workers=opts['workers']
)
encoder_loss = EncoderLoss()
discriminator_grad_reg = DiscGradReg()
# Load Generator in evaluation mode
with dnnlib.util.open_url(opts['pkl']) as f:
G = legacy.load_network_pkl(f)['G_ema'].to(device)
G.eval()
with dnnlib.util.open_url(opts['pkl']) as f:
D = legacy.load_network_pkl(f)['D'].to(device)
D.train()
E = FastGANEncoder()
E.to(device)
E.train()
# Initialize optimizers
optim_d = AdamW(D.parameters())
optim_e = AdamW(E.parameters())
print(f"\n\nTraining for {opts['epochs']} epochs with batch size {str(opts['batch_gpu']) + ' and total batch ' + str(opts['batch']) if opts.get('batch_gpu') else int(opts['batch'])} on {num_classes} classes...\n\n")
for epoch in range(opts['epochs']):
running_loss_d = 0.
running_loss_e = 0.
# Iterate over dataset
iters = 0
for i, (imgs, labels) in enumerate(dataloader, 0):
with autocast():
imgs, labels = imgs.cuda(), labels.cuda()
# Encode the batch
z_pred = E(imgs)
# Pass the batch through the generator
reconsts = G(z_pred, labels)
# Pass the images through the discriminator
fake_score = D(reconsts, labels)
real_score = D(imgs.requires_grad_(), labels)
# Calculate the loss
loss_e = encoder_loss(imgs, reconsts, fake_score)
loss_d = fake_score - real_score + discriminator_grad_reg(imgs, real_score.requires_grad_())
# Determine whether or not to back prop
if (iters + 1) % accumulate_steps == 0:
iters = 0
loss_e.backward()
loss_d.backward()
optim_e.step()
optim_d.step()
iters += 1
print(f"Running loss of the discriminator at epoch {epoch + 1}: {running_loss_d}")
print(f"Running loss of the encoder at epoch {epoch + 1}: {running_loss_e}")
I hope that helps, I truly do appreciate your help. The function interfaces with the click library, so that’s why it loads in options from the dictionary opts.