class BiGAN(object):
def __init__(self, args):
self.z_dim = args.z_dim
self.decay_rate = args.decay_rate
self.learning_rate = args.learning_rate
self.model_name = args.model_name
self.batch_size = args.batch_size
#initialize networks
self.Generator = Generator(self.z_dim).cuda()
self.Encoder = Encoder(self.z_dim).cuda()
self.Discriminator = Discriminator().cuda()
#set optimizers for all networks
self.optimizer_G_E = torch.optim.Adam(
list(self.Generator.parameters()) +
list(self.Encoder.parameters()),
lr=self.learning_rate,
betas=(0.5, 0.999))
self.optimizer_D = torch.optim.Adam(self.Discriminator.parameters(),
lr=self.learning_rate,
betas=(0.5, 0.999))
#initialize network weights
self.Generator.apply(weights_init)
self.Encoder.apply(weights_init)
self.Discriminator.apply(weights_init)
def train(self, data):
self.Generator.train()
self.Encoder.train()
self.Discriminator.train()
self.optimizer_G_E.zero_grad()
self.optimizer_D.zero_grad()
#get fake z_data for generator
self.z_fake = torch.randn((self.batch_size, self.z_dim))
#send fake z_data through generator to get fake x_data
self.x_fake = self.Generator(self.z_fake.detach())
#send real data through encoder to get real z_data
self.z_real = self.Encoder(data)
#send real x and z data into discriminator
self.out_real = self.Discriminator(data, z_real.detach())
#send fake x and z data into discriminator
self.out_fake = self.Discriminator(x_fake.detach(), z_fake.detach())
#compute discriminator loss
self.D_loss = nn.BCELoss()
#compute generator/encoder loss
self.G_E_loss = nn.BCELoss()
#compute discriminator gradiants and backpropogate
self.D_loss.backward()
self.optimizer_D.step()
#compute generator/encoder gradiants and backpropogate
self.G_E_loss.backward()
self.optimizer_G_E.step()
opt = parser.parse_args()
print(opt)
img_shape = (opt.channels, opt.img_size, opt.img_size)
generator = Generator(dim = 64, zdim=opt.latent_dim, nc=opt.channels)
discriminator = Discriminator(dim = 64, zdim=opt.latent_dim, nc=opt.channels,out_feat=True)
encoder = Encoder(dim = 64, zdim=opt.latent_dim, nc=opt.channels)
generator.load_state_dict(torch.load(opt.G_path))
discriminator.load_state_dict(torch.load(opt.D_path))
generator.to(opt.device)
encoder.to(opt.device)
discriminator.to(opt.device)
encoder.train()
discriminator.train()
dataloader = load_data(opt)
generator.eval()
Tensor = torch.cuda.FloatTensor if opt.device == 'cuda' else torch.FloatTensor
optimizer_E = torch.optim.Adam(encoder.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
optimizer_D = torch.optim.Adam(discriminator.parameters(), lr=opt.lr, betas=(opt.b1, opt.b2))
max_auc = 0
for epoch in range(opt.n_epochs):
# train
# Criterion
l1_norm = nn.L1Loss()
it = 0
decayed_lr_G = args.lr_G
decayed_lr_D = args.lr_D
total_epochs = args.epochs + args.epochs_decay
for ep in range(start_ep, total_epochs):
# Linearly decay learning rates
if ep >= args.epochs:
decayed_lr_G = args.lr_G / args.epochs_decay * (total_epochs - ep)
decayed_lr_D = args.lr_D / args.epochs_decay * (total_epochs - ep)
set_lr(G_opt, decayed_lr_G)
set_lr(D_opt, decayed_lr_D)
# Optimize parameters
E.train()
G.train()
D.train()
for reals, annos in tqdm(train_data):
reals, annos = reals.to(device), annos.to(device)
annos_onehot = onehot2d(annos, n_classes).type_as(reals)
# Train D
trainable(E, False)
trainable(G, False)
trainable(D, True)
mu, logvar = E(reals)
latents = sample_latent(mu, logvar).detach()
fakes = G(latents, annos_onehot).detach()
d_real = D(reals, annos_onehot)
d_fake = D(fakes, annos_onehot)
# Real/fake hinge loss