def testSyncBatchNormSyncEval(self):
bn = nn.BatchNorm1d(10, eps=1e-5, affine=False)
sync_bn = SynchronizedBatchNorm1d(10, eps=1e-5, affine=False)
sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
bn.cuda()
sync_bn.cuda()
self._checkBatchNormResult(bn,
sync_bn,
torch.rand(16, 10),
False,
cuda=True)
def testSyncBatchNorm2DSyncTrain(self):
bn = nn.BatchNorm2d(10)
sync_bn = SynchronizedBatchNorm2d(10)
sync_bn = DataParallelWithCallback(sync_bn, device_ids=[0, 1])
bn.cuda()
sync_bn.cuda()
self._checkBatchNormResult(bn,
sync_bn,
torch.rand(16, 10, 16, 16),
True,
cuda=True)
def main(args):
# parameters
img_size = args.img_size
z_dim = 128
lamb_obj = 1.0
lamb_img = 0.1
num_classes = 184 if args.dataset == 'coco' else 179
num_obj = 8 if args.dataset == 'coco' else 31
args.out_path = os.path.join(args.out_path, args.dataset,
str(args.img_size))
num_gpus = torch.cuda.device_count()
num_workers = 2
if num_gpus > 1:
parallel = True
args.batch_size = args.batch_size * num_gpus
num_workers = num_workers * num_gpus
else:
parallel = False
# data loader
train_data = get_dataset(args.dataset, img_size)
dataloader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
num_workers=8)
# Load model
device = torch.device('cuda')
netG = context_aware_generator(num_classes=num_classes,
output_dim=3).to(device)
netD = CombineDiscriminator128(num_classes=num_classes).to(device)
parallel = True
if parallel:
netG = DataParallelWithCallback(netG)
netD = nn.DataParallel(netD)
g_lr, d_lr = args.g_lr, args.d_lr
gen_parameters = []
for key, value in dict(netG.named_parameters()).items():
if value.requires_grad:
if 'mapping' in key:
gen_parameters += [{'params': [value], 'lr': g_lr * 0.1}]
else:
gen_parameters += [{'params': [value], 'lr': g_lr}]
g_optimizer = torch.optim.Adam(gen_parameters, betas=(0, 0.999))
dis_parameters = []
for key, value in dict(netD.named_parameters()).items():
if value.requires_grad:
dis_parameters += [{'params': [value], 'lr': d_lr}]
d_optimizer = torch.optim.Adam(dis_parameters, betas=(0, 0.999))
# make dirs
if not os.path.exists(args.out_path):
os.makedirs(args.out_path)
if not os.path.exists(os.path.join(args.out_path, 'model/')):
os.makedirs(os.path.join(args.out_path, 'model/'))
writer = SummaryWriter(os.path.join(args.out_path, 'log'))
logger = setup_logger("lostGAN", args.out_path, 0)
logger.info(netG)
logger.info(netD)
start_time = time.time()
vgg_loss = VGGLoss()
vgg_loss = nn.DataParallel(vgg_loss)
l1_loss = nn.DataParallel(nn.L1Loss())
for epoch in range(args.total_epoch):
netG.train()
netD.train()
for idx, data in enumerate(tqdm(dataloader)):
real_images, label, bbox = data
# print(real_images.shape)
# print(label.shape)
# print(bbox.shape)
real_images, label, bbox = real_images.to(device), label.long().to(
device).unsqueeze(-1), bbox.float()
# update D network
netD.zero_grad()
real_images, label = real_images.to(device), label.long().to(
device)
d_out_real, d_out_robj = netD(real_images, bbox, label)
d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()
d_loss_robj = torch.nn.ReLU()(1.0 - d_out_robj).mean()
z = torch.randn(real_images.size(0), num_obj, z_dim).to(device)
fake_images = netG(z, bbox, y=label.squeeze(dim=-1))
d_out_fake, d_out_fobj = netD(fake_images.detach(), bbox, label)
d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()
d_loss_fobj = torch.nn.ReLU()(1.0 + d_out_fobj).mean()
d_loss = lamb_obj * (d_loss_robj + d_loss_fobj) + lamb_img * (
d_loss_real + d_loss_fake)
d_loss.backward()
d_optimizer.step()
# update G network
if (idx % 1) == 0:
netG.zero_grad()
g_out_fake, g_out_obj = netD(fake_images, bbox, label)
g_loss_fake = -g_out_fake.mean()
g_loss_obj = -g_out_obj.mean()
pixel_loss = l1_loss(fake_images, real_images).mean()
feat_loss = vgg_loss(fake_images, real_images).mean()
g_loss = g_loss_obj * lamb_obj + g_loss_fake * lamb_img + pixel_loss + feat_loss
g_loss.backward()
g_optimizer.step()
if (idx + 1) % 10 == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
logger.info("Time Elapsed: [{}]".format(elapsed))
logger.info(
"Step[{}/{}], d_out_real: {:.4f}, d_out_fake: {:.4f}, g_out_fake: {:.4f} "
.format(epoch + 1, idx + 1, d_loss_real.item(),
d_loss_fake.item(), g_loss_fake.item()))
logger.info(
" d_obj_real: {:.4f}, d_obj_fake: {:.4f}, g_obj_fake: {:.4f} "
.format(d_loss_robj.item(), d_loss_fobj.item(),
g_loss_obj.item()))
logger.info(
" pixel_loss: {:.4f}, feat_loss: {:.4f}".
format(pixel_loss.item(), feat_loss.item()))
writer.add_image(
"real images",
make_grid(real_images.cpu().data * 0.5 + 0.5, nrow=4),
epoch * len(dataloader) + idx + 1)
writer.add_image(
"fake images",
make_grid(fake_images.cpu().data * 0.5 + 0.5, nrow=4),
epoch * len(dataloader) + idx + 1)
writer.add_scalars(
"D_loss_real", {
"real": d_loss_real.item(),
"robj": d_loss_robj.item(),
"loss": d_loss.item()
})
writer.add_scalars("D_loss_fake", {
"fake": d_loss_fake.item(),
"fobj": d_loss_fobj.item()
})
writer.add_scalars(
"G_loss", {
"fake": g_loss_fake.item(),
"obj": g_loss_obj.item(),
"loss": g_loss.item()
})
# save model
if (epoch + 1) % 5 == 0:
torch.save(
netG.state_dict(),
os.path.join(args.out_path, 'model/',
'G_%d.pth' % (epoch + 1)))
def main(args):
# parameters
img_size = 128
z_dim = 128
lamb_obj = 1.0
lamb_app = 1.0
lamb_img = 0.1
num_classes = 184 if args.dataset == 'coco' else 179
num_obj = 8 if args.dataset == 'coco' else 31
# args.out_path = os.path.join(args.out_path, args.dataset + '_1gpu', str(img_size))
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_ids
# data loader
train_data = get_dataset(args.dataset, img_size)
print(f"{args.dataset.title()} datasets with {len(train_data)} samples has been created!")
num_gpus = torch.cuda.device_count()
num_workers = 20
if num_gpus > 1:
parallel = True
args.batch_size = args.batch_size * num_gpus
num_workers = num_workers * num_gpus
else:
parallel = False
print("{} GPUs, {} workers are used".format(num_gpus, num_workers))
dataloader = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size,
drop_last=True, shuffle=True, num_workers=num_workers)
# Load model
device = torch.device('cuda')
netG = ResnetGenerator128(num_classes=num_classes, output_dim=3).to(device)
netD = CombineDiscriminator128_app(num_classes=num_classes).to(device)
# if os.path.isfile(args.checkpoint):
# state_dict = torch.load(args.checkpoint)
# parallel = True
if parallel:
netG = DataParallelWithCallback(netG)
netD = nn.DataParallel(netD)
g_lr, d_lr = args.g_lr, args.d_lr
gen_parameters = []
for key, value in dict(netG.named_parameters()).items():
if value.requires_grad:
if 'mapping' in key:
gen_parameters += [{'params': [value], 'lr': g_lr * 0.1}]
else:
gen_parameters += [{'params': [value], 'lr': g_lr}]
g_optimizer = torch.optim.Adam(gen_parameters, betas=(0, 0.999))
dis_parameters = []
for key, value in dict(netD.named_parameters()).items():
if value.requires_grad:
dis_parameters += [{'params': [value], 'lr': d_lr}]
d_optimizer = torch.optim.Adam(dis_parameters, betas=(0, 0.999))
# make dirs
if not os.path.exists(args.out_path):
os.makedirs(args.out_path)
if not os.path.exists(os.path.join(args.out_path, 'model/')):
os.makedirs(os.path.join(args.out_path, 'model/'))
writer = SummaryWriter(os.path.join(args.out_path, 'log'))
# writer = None
logger = setup_logger("lostGAN", args.out_path, 0)
for arg in vars(args):
logger.info("%15s : %-15s" % (arg, str(getattr(args, arg))))
logger.info(netG)
logger.info(netD)
start_time = time.time()
vgg_loss = VGGLoss()
vgg_loss = nn.DataParallel(vgg_loss)
l1_loss = nn.DataParallel(nn.L1Loss())
for epoch in range(args.total_epoch):
netG.train()
netD.train()
print("Epoch {}/{}".format(epoch, args.total_epoch))
for idx, data in enumerate(tqdm(dataloader)):
real_images, label, bbox = data
real_images, label, bbox = real_images.to(device), label.long().to(device).unsqueeze(-1), bbox.float()
# update D network
netD.zero_grad()
real_images, label = real_images.to(device), label.long().to(device)
d_out_real, d_out_robj, d_out_robj_app = netD(real_images, bbox, label)
d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()
d_loss_robj = torch.nn.ReLU()(1.0 - d_out_robj).mean()
d_loss_robj_app = torch.nn.ReLU()(1.0 - d_out_robj_app).mean()
# print(d_loss_robj)
# print(d_loss_robj_app)
z = torch.randn(real_images.size(0), num_obj, z_dim).to(device)
fake_images = netG(z, bbox, y=label.squeeze(dim=-1))
d_out_fake, d_out_fobj, d_out_fobj_app = netD(fake_images.detach(), bbox, label)
d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()
d_loss_fobj = torch.nn.ReLU()(1.0 + d_out_fobj).mean()
d_loss_fobj_app = torch.nn.ReLU()(1.0 + d_out_fobj_app).mean()
d_loss = lamb_obj * (d_loss_robj + d_loss_fobj) + lamb_img * (d_loss_real + d_loss_fake) + lamb_app * (d_loss_robj_app + d_loss_fobj_app)
d_loss.backward()
d_optimizer.step()
# update G network
if (idx % 1) == 0:
netG.zero_grad()
g_out_fake, g_out_obj, g_out_obj_app = netD(fake_images, bbox, label)
g_loss_fake = - g_out_fake.mean()
g_loss_obj = - g_out_obj.mean()
g_loss_obj_app = - g_out_obj_app.mean()
pixel_loss = l1_loss(fake_images, real_images).mean()
feat_loss = vgg_loss(fake_images, real_images).mean()
g_loss = g_loss_obj * lamb_obj + g_loss_fake * lamb_img + pixel_loss + feat_loss + lamb_app * g_loss_obj_app
g_loss.backward()
g_optimizer.step()
# print("d_loss_real={:.3f}, d_loss_robj={:.3f}, d_loss_robj_app={:.3f}".format(d_loss_real.item(), d_loss_robj.item(), d_loss_robj_app.item()))
# print("d_loss_fake={:.3f}, d_loss_fobj={:.3f}, d_loss_fobj_app={:.3f}".format(d_loss_fake.item(), d_loss_fobj.item(), d_loss_fobj_app.item()))
# print("g_loss_fake={:.3f}, g_loss_obj={:.3f}, g_loss_obj_app={:.3f}".format(g_loss_fake.item(), g_loss_obj.item(), g_loss_obj_app.item()))
if (idx + 1) % 100 == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
logger.info("Time Elapsed: [{}]".format(elapsed))
logger.info("Step[{}/{}], d_out_real: {:.4f}, d_out_fake: {:.4f}, g_out_fake: {:.4f} ".format(epoch + 1,
idx + 1,
d_loss_real.item(),
d_loss_fake.item(),
g_loss_fake.item()))
logger.info(" d_obj_real: {:.4f}, d_obj_fake: {:.4f}, g_obj_fake: {:.4f} ".format(
d_loss_robj.item(),
d_loss_fobj.item(),
g_loss_obj.item()))
logger.info(" d_obj_real_app: {:.4f}, d_obj_fake_app: {:.4f}, g_obj_fake_app: {:.4f} ".format(
d_loss_robj_app.item(),
d_loss_fobj_app.item(),
g_loss_obj_app.item()))
logger.info(" pixel_loss: {:.4f}, feat_loss: {:.4f}".format(pixel_loss.item(), feat_loss.item()))
if writer is not None:
writer.add_image("real images", make_grid(real_images.cpu().data * 0.5 + 0.5, nrow=4), epoch * len(dataloader) + idx + 1)
writer.add_image("fake images", make_grid(fake_images.cpu().data * 0.5 + 0.5, nrow=4), epoch * len(dataloader) + idx + 1)
writer.add_scalars("D_loss_real", {"real": d_loss_real.item(),
"robj": d_loss_robj.item(),
"robj_app": d_loss_robj_app.item(),
"loss": d_loss.item()})
writer.add_scalars("D_loss_fake", {"fake": d_loss_fake.item(),
"fobj": d_loss_fobj.item(),
"fobj_app": d_loss_fobj_app.item()})
writer.add_scalars("G_loss", {"fake": g_loss_fake.item(),
"obj_app": g_loss_obj_app.item(),
"obj": g_loss_obj.item(),
"loss": g_loss.item()})
# save model
if (epoch + 1) % 5 == 0:
torch.save(netG.state_dict(), os.path.join(args.out_path, 'model/', 'G_%d.pth' % (epoch + 1)))
torch.save(netD.state_dict(), os.path.join(args.out_path, 'model/', 'D_%d.pth' % (epoch + 1)))
def main(args):
# parameters
img_size = 128
z_dim = 128
lamb_obj = 1.1 # 0.5, 1, 1, 0
lamb_img = 0.1 # 1, 1, 0.5, 1
num_classes = 184 if args.dataset == 'coco' else 10 # 179
num_obj = 8 if args.dataset == 'coco' else 9 # 31
# data loader
train_data = get_dataset(args.dataset, img_size)
dataloader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
num_workers=8)
# Load model
netG = ResnetGenerator128(num_classes=num_classes, output_dim=3).cuda()
netD = CombineDiscriminator128(num_classes=num_classes).cuda()
parallel = True
if parallel:
netG = DataParallelWithCallback(netG)
netD = nn.DataParallel(netD)
g_lr, d_lr = args.g_lr, args.d_lr
gen_parameters = []
for key, value in dict(netG.named_parameters()).items():
if value.requires_grad:
if 'mapping' in key:
gen_parameters += [{'params': [value], 'lr': g_lr * 0.1}]
else:
gen_parameters += [{'params': [value], 'lr': g_lr}]
g_optimizer = torch.optim.Adam(gen_parameters, betas=(0, 0.999))
dis_parameters = []
for key, value in dict(netD.named_parameters()).items():
if value.requires_grad:
dis_parameters += [{'params': [value], 'lr': d_lr}]
d_optimizer = torch.optim.Adam(dis_parameters, betas=(0, 0.999))
# get most recent commit
commit_obj = get_recent_commit()
current_time = time.strftime("%H_%M_%dd_%mm", time.localtime())
# make dirs
if not os.path.exists(args.out_path):
os.makedirs(args.out_path)
if not os.path.exists(os.path.join(args.out_path, 'model/')):
os.makedirs(os.path.join(args.out_path, 'model/'))
writer = SummaryWriter(os.path.join(args.out_path, 'log'))
logger = setup_logger("lostGAN",
args.out_path,
0,
filename=current_time + '_log.txt')
logger.info('Commit Tag: ' + commit_obj)
logger.info('Time: ', time.localtime())
logger.info('No rotations. No resizing. Ob=1.1, im=0.1. 3x3 grid.')
logger.info(netG)
logger.info(netD)
# labelencoder = LabelEncoder()
# enc = OneHotEncoder()
total_steps = len(dataloader)
start_time = time.time()
for epoch in range(args.total_epoch):
netG.train()
netD.train()
print('epoch ', epoch)
for idx, data in enumerate(dataloader):
real_images, label, bbox = data
real_images, label, bbox = real_images.cuda(), \
label.long().cuda().unsqueeze(-1), \
bbox.float() # , \
# rotation.float() # .unsqueeze(-1)
# rotation = labelencoder.fit_transform(rotation.view(-1)).reshape(-1, 16)
# rotation = torch.from_numpy(rotation).float().cuda()
# rotation = enc.fit_transform(rotation).toarray()
# rotation = torch.from_numpy(rotation).float().cuda() # [bs*16, 4]
# update D network
netD.zero_grad()
real_images, label = real_images.float().cuda(), label.long().cuda(
)
d_out_real, d_out_robj = netD(real_images, bbox, label)
# rotation=rotation) # d_out_robj: [460, 1]
d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()
d_loss_robj = torch.nn.ReLU()(1.0 - d_out_robj).mean()
# d_loss_rot = torch.nn.ReLU()(1.0 - d_out_rot).mean()
z = torch.randn(real_images.size(0), num_obj, z_dim).cuda()
fake_images = netG(z, bbox, y=label.squeeze(dim=-1)) # rm rot
d_out_fake, d_out_fobj = netD(fake_images.detach(), bbox,
label) # rm rot
d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()
d_loss_fobj = torch.nn.ReLU()(1.0 + d_out_fobj).mean()
# d_loss_frot = torch.nn.ReLU()(1.0 + d_out_frot).mean()
d_loss = lamb_obj * (d_loss_robj + d_loss_fobj) + \
lamb_img * (d_loss_real + d_loss_fake)
d_loss.backward()
d_optimizer.step()
# update G network
if (idx % 1) == 0:
netG.zero_grad()
g_out_fake, g_out_obj = netD(fake_images, bbox,
label) # rm rot
g_loss_fake = -g_out_fake.mean()
g_loss_obj = -g_out_obj.mean()
# g_loss_rot = - g_out_rot.mean()
g_loss = g_loss_obj * lamb_obj + g_loss_fake * lamb_img
g_loss.backward()
g_optimizer.step()
if (idx + 1) % 100 == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
logger.info("Time Elapsed: [{}]".format(elapsed))
logger.info(
"Epoch[{}/{}], Step[{}/{}], "
"d_out_real: {:.4f}, d_out_fake: {:.4f}, g_out_fake: {:.4f} "
.format(epoch + 1, args.total_epoch, idx + 1, total_steps,
d_loss_real.item(), d_loss_fake.item(),
g_loss_fake.item()))
logger.info(
"d_obj_real: {:.4f}, d_obj_fake: {:.4f}, g_obj_fake: {:.4f} "
.format(d_loss_robj.item(), d_loss_fobj.item(),
g_loss_obj.item()))
writer.add_image(
"real images",
make_grid(real_images.cpu().data * 0.5 + 0.5, nrow=4),
epoch * len(dataloader) + idx + 1)
writer.add_image(
"fake images",
make_grid(fake_images.cpu().data * 0.5 + 0.5, nrow=4),
epoch * len(dataloader) + idx + 1)
writer.add_scalar('DLoss', d_loss,
epoch * len(dataloader) + idx + 1)
writer.add_scalar('DLoss/real_images', d_loss_real,
epoch * len(dataloader) + idx + 1)
writer.add_scalar('DLoss/fake_images', d_loss_fake,
epoch * len(dataloader) + idx + 1)
writer.add_scalar('DLoss/real_objects', d_loss_robj,
epoch * len(dataloader) + idx + 1)
writer.add_scalar('DLoss/fake_objects', d_loss_fobj,
epoch * len(dataloader) + idx + 1)
writer.add_scalar('GLoss', g_loss,
epoch * len(dataloader) + idx + 1)
writer.add_scalar('GLoss/fake_images',
epoch * len(dataloader) + idx + 1)
writer.add_scalar('GLoss/fake_objects',
epoch * len(dataloader) + idx + 1)
# save model
if (epoch + 1) % 5 == 0:
torch.save(
netG.state_dict(),
os.path.join(args.out_path, 'model/',
'G_%d.pth' % (epoch + 1)))
def main(args):
# parameters
img_size = 128
z_dim = 128
lamb_obj = 1.0
lamb_app = 1.0
lamb_img = 0.1
num_classes = 184 if args.dataset == 'coco' else 179
num_obj = 8 if args.dataset == 'coco' else 31
args.out_path = os.path.join(args.out_path, args.dataset, str(img_size))
num_gpus = torch.cuda.device_count()
num_workers = 2
if num_gpus > 1:
parallel = True
args.batch_size = args.batch_size * num_gpus
num_workers = num_workers * num_gpus
else:
parallel = False
# data loader
train_data = get_dataset(args.dataset, img_size)
dataloader = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
num_workers=num_workers)
# Load model
device = torch.device('cuda')
netG = ResnetGenerator128_context(num_classes=num_classes,
output_dim=3).to(device)
netD = CombineDiscriminator128_app(num_classes=num_classes).to(device)
if args.checkpoint_epoch is not None:
load_G = args.out_path + '/model/G_{}.pth'.format(
args.checkpoint_epoch)
load_D = args.out_path + '/model/D_{}.pth'.format(
args.checkpoint_epoch)
if not os.path.isfile(load_G) or not os.path.isfile(load_D):
return
# load generator
state_dict = torch.load(load_G)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`nvidia
new_state_dict[name] = v
model_dict = netG.state_dict()
pretrained_dict = {
k: v
for k, v in new_state_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
netG.load_state_dict(model_dict)
netG.cuda()
# load discriminator
state_dict = torch.load(load_D)
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove `module.`nvidia
new_state_dict[name] = v
model_dict = netD.state_dict()
pretrained_dict = {
k: v
for k, v in new_state_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
netD.load_state_dict(model_dict)
netD.cuda()
else:
args.checkpoint_epoch = 0
# parallel = False
if parallel:
netG = DataParallelWithCallback(netG)
netD = nn.DataParallel(netD)
g_lr, d_lr = args.g_lr, args.d_lr
gen_parameters = []
for key, value in dict(netG.named_parameters()).items():
if value.requires_grad:
if 'mapping' in key:
gen_parameters += [{'params': [value], 'lr': g_lr * 0.1}]
else:
gen_parameters += [{'params': [value], 'lr': g_lr}]
g_optimizer = torch.optim.Adam(gen_parameters, betas=(0, 0.999))
dis_parameters = []
for key, value in dict(netD.named_parameters()).items():
if value.requires_grad:
dis_parameters += [{'params': [value], 'lr': d_lr}]
d_optimizer = torch.optim.Adam(dis_parameters, betas=(0, 0.999))
# make dirs
if not os.path.exists(args.out_path):
os.makedirs(args.out_path)
if not os.path.exists(os.path.join(args.out_path, 'model/')):
os.makedirs(os.path.join(args.out_path, 'model/'))
writer = SummaryWriter(os.path.join(args.out_path, 'log'))
logger = setup_logger("lostGAN", args.out_path, 0)
logger.info(netG)
logger.info(netD)
start_time = time.time()
vgg_loss = VGGLoss()
vgg_loss = nn.DataParallel(vgg_loss)
l1_loss = nn.DataParallel(nn.L1Loss())
for epoch in range(args.checkpoint_epoch, args.total_epoch):
netG.train()
netD.train()
print("Epoch {}/{}".format(epoch, args.total_epoch))
for idx, data in enumerate(tqdm(dataloader)):
real_images, label, bbox = data
# print(real_images.shape)
# print(label.shape)
# print(bbox.shape)
real_images, label, bbox = real_images.to(device), label.long().to(
device).unsqueeze(-1), bbox.float()
# update D network
netD.zero_grad()
real_images, label = real_images.to(device), label.long().to(
device)
d_out_real, d_out_robj, d_out_robj_app = netD(
real_images, bbox, label)
d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()
d_loss_robj = torch.nn.ReLU()(1.0 - d_out_robj).mean()
d_loss_robj_app = torch.nn.ReLU()(1.0 - d_out_robj_app).mean()
# print(d_loss_robj)
# print(d_loss_robj_app)
z = torch.randn(real_images.size(0), num_obj, z_dim).to(device)
fake_images = netG(z, bbox, y=label.squeeze(dim=-1))
d_out_fake, d_out_fobj, d_out_fobj_app = netD(
fake_images.detach(), bbox, label)
d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()
d_loss_fobj = torch.nn.ReLU()(1.0 + d_out_fobj).mean()
d_loss_fobj_app = torch.nn.ReLU()(1.0 + d_out_fobj_app).mean()
d_loss = lamb_obj * (d_loss_robj + d_loss_fobj) + lamb_img * (
d_loss_real + d_loss_fake) + lamb_app * (d_loss_robj_app +
d_loss_fobj_app)
d_loss.backward()
d_optimizer.step()
# update G network
if (idx % 1) == 0:
netG.zero_grad()
g_out_fake, g_out_obj, g_out_obj_app = netD(
fake_images, bbox, label)
g_loss_fake = -g_out_fake.mean()
g_loss_obj = -g_out_obj.mean()
g_loss_obj_app = -g_out_obj_app.mean()
pixel_loss = l1_loss(fake_images, real_images).mean()
feat_loss = vgg_loss(fake_images, real_images).mean()
g_loss = g_loss_obj * lamb_obj + g_loss_fake * lamb_img + pixel_loss + feat_loss + lamb_app * g_loss_obj_app
g_loss.backward()
g_optimizer.step()
if (idx + 1) % 500 == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
logger.info("Time Elapsed: [{}]".format(elapsed))
logger.info(
"Step[{}/{}], d_out_real: {:.4f}, d_out_fake: {:.4f}, g_out_fake: {:.4f} "
.format(epoch + 1, idx + 1, d_loss_real.item(),
d_loss_fake.item(), g_loss_fake.item()))
logger.info(
" d_obj_real: {:.4f}, d_obj_fake: {:.4f}, g_obj_fake: {:.4f} "
.format(d_loss_robj.item(), d_loss_fobj.item(),
g_loss_obj.item()))
logger.info(
" d_obj_real_app: {:.4f}, d_obj_fake_app: {:.4f}, g_obj_fake_app: {:.4f} "
.format(d_loss_robj_app.item(), d_loss_fobj_app.item(),
g_loss_obj_app.item()))
logger.info(
" pixel_loss: {:.4f}, feat_loss: {:.4f}".
format(pixel_loss.item(), feat_loss.item()))
writer.add_image(
"real images",
make_grid(real_images.cpu().data * 0.5 + 0.5, nrow=4),
epoch * len(dataloader) + idx + 1)
writer.add_image(
"fake images",
make_grid(fake_images.cpu().data * 0.5 + 0.5, nrow=4),
epoch * len(dataloader) + idx + 1)
# save model
if (epoch + 1) % 5 == 0:
torch.save(
netG.state_dict(),
os.path.join(args.out_path, 'model/',
'G_%d.pth' % (epoch + 1)))
torch.save(
netD.state_dict(),
os.path.join(args.out_path, 'model/',
'D_%d.pth' % (epoch + 1)))