def main():
dataset = MyImageFolder(root_dir="data/")
loader = DataLoader(
dataset,
batch_size=config.BATCH_SIZE,
shuffle=True,
pin_memory=True,
num_workers=config.NUM_WORKERS,
)
gen = Generator(in_channels=3).to(config.DEVICE)
disc = Discriminator(in_channels=3).to(config.DEVICE)
initialize_weights(gen)
opt_gen = optim.Adam(gen.parameters(),
lr=config.LEARNING_RATE,
betas=(0.0, 0.9))
opt_disc = optim.Adam(disc.parameters(),
lr=config.LEARNING_RATE,
betas=(0.0, 0.9))
writer = SummaryWriter("logs")
tb_step = 0
l1 = nn.L1Loss()
gen.train()
disc.train()
vgg_loss = VGGLoss()
g_scaler = torch.cuda.amp.GradScaler()
d_scaler = torch.cuda.amp.GradScaler()
if config.LOAD_MODEL:
load_checkpoint(
config.CHECKPOINT_GEN,
gen,
opt_gen,
config.LEARNING_RATE,
)
load_checkpoint(
config.CHECKPOINT_DISC,
disc,
opt_disc,
config.LEARNING_RATE,
)
for epoch in range(config.NUM_EPOCHS):
tb_step = train_fn(
loader,
disc,
gen,
opt_gen,
opt_disc,
l1,
vgg_loss,
g_scaler,
d_scaler,
writer,
tb_step,
)
if config.SAVE_MODEL:
save_checkpoint(gen, opt_gen, filename=config.CHECKPOINT_GEN)
save_checkpoint(disc, opt_disc, filename=config.CHECKPOINT_DISC)
def build_model(self, config):
"""Create a generator and a discriminator."""
self.G = Generator(config.g_conv_dim, config.d_channel,
config.channel_1x1) # 2 for mask vector.
self.D = Discriminator(config.crop_size, config.d_conv_dim,
config.d_repeat_num)
self.G.apply(weights_init)
self.D.apply(weights_init)
self.G.cuda()
self.D.cuda()
self.g_optimizer = torch.optim.Adam(self.G.parameters(), config.g_lr,
[config.beta1, config.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), config.d_lr,
[config.beta1, config.beta2])
self.G = nn.DataParallel(self.G)
self.D = nn.DataParallel(self.D)
self.VGGLoss = VGGLoss().eval()
self.VGGLoss.cuda()
self.VGGLoss = nn.DataParallel(self.VGGLoss)
def main():
dataset = MyImageFolder(root_dir="new_data/")
loader = DataLoader(
dataset,
batch_size=config.BATCH_SIZE,
shuffle=True,
pin_memory=True,
num_workers=config.NUM_WORKERS,
)
gen = Generator(in_channels=3).to(config.DEVICE)
disc = Discriminator(img_channels=3).to(config.DEVICE)
opt_gen = optim.Adam(gen.parameters(),
lr=config.LEARNING_RATE,
betas=(0.9, 0.999))
opt_disc = optim.Adam(disc.parameters(),
lr=config.LEARNING_RATE,
betas=(0.9, 0.999))
mse = nn.MSELoss()
bce = nn.BCEWithLogitsLoss()
vgg_loss = VGGLoss()
if config.LOAD_MODEL:
load_checkpoint(
config.CHECKPOINT_GEN,
gen,
opt_gen,
config.LEARNING_RATE,
)
load_checkpoint(
config.CHECKPOINT_DISC,
disc,
opt_disc,
config.LEARNING_RATE,
)
for epoch in range(config.NUM_EPOCHS):
train_fn(loader, disc, gen, opt_gen, opt_disc, mse, bce, vgg_loss)
if config.SAVE_MODEL:
save_checkpoint(gen, opt_gen, filename=config.CHECKPOINT_GEN)
save_checkpoint(disc, opt_disc, filename=config.CHECKPOINT_DISC)
def main():
dataset = ImageDS(root="Data/")
loader = DataLoader(
dataset,
batch_size=config.BATCH_SIZE,
shuffle=True,
pin_memory=True,
)
gen = Generator().to(config.DEVICE)
disc = Discriminator().to(config.DEVICE)
opt_gen = optim.Adam(gen.parameters(),
lr=config.LEARNING_RATE,
betas=(0.9, 0.999))
opt_disc = optim.Adam(disc.parameters(),
lr=config.LEARNING_RATE,
betas=(0.9, 0.999))
mse = nn.MSELoss()
bce = nn.BCEWithLogitsLoss()
vgg_loss = VGGLoss()
for epoch in range(config.NUM_EPOCHS):
train_fn(loader, disc, gen, opt_gen, opt_disc, mse, bce, vgg_loss)
batch_size=args.batch_size, \
shuffle=False, \
num_workers=0)
editor_eval = Editor(solo, reconstructor)
editor_eval.load_state_dict(torch.load(args.PATH))
eval(editor_eval.to(device), coco_test_loader)
else:
if not os.path.exists('checkpoints/'):
os.makedirs('checkpoints/')
logger.info("Instantiating Editor")
editor = Editor(solo, reconstructor)
coco_train_loader, _ = get_loader(device=device, \
# root=args.coco+'train2017', \
root = 'datasets/bg_composite/', \
json=args.coco+'annotations/instances_train2017.json', \
batch_size=args.batch_size, \
shuffle=True, \
num_workers=0)
if args.load:
editor.load_state_dict(torch.load(args.PATH))
editor.to(device)
vgg_loss = VGGLoss()
recon_loss = ReconLoss()
train(model=editor,
num_epochs=args.num_epochs,
dataloader=coco_train_loader)
#ssim0 = 1 - ssim_loss(d0,labels_v)
# iou0 = iou_loss(d0,labels_v)
#loss = torch.pow(torch.mean(torch.abs(labels_v-d0)),2)*(5.0*loss0 + loss1 + loss2 + loss3 + loss4 + loss5) #+ 5.0*lossa
loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6 + loss7 #+ 5.0*lossa
print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n" %
(loss0.item(), loss1.item(), loss2.item(), loss3.item(),
loss4.item(), loss5.item(), loss6.item()))
# print("BCE: l1:%3f, l2:%3f, l3:%3f, l4:%3f, l5:%3f, la:%3f, all:%3f\n"%(loss1.data[0],loss2.data[0],loss3.data[0],loss4.data[0],loss5.data[0],lossa.data[0],loss.data[0]))
return loss
criterion_GAN = torch.nn.MSELoss().cuda()
criterion_bce = torch.nn.BCEWithLogitsLoss().cuda()
criterion_vgg = VGGLoss().cuda()
gen1 = BASNet()
gen2 = ressenet()
writer = SummaryWriter(
log_dir="/public/zebanghe2/derain/reimplement/residualse/modeltest/",
comment="ResSENet")
gen1 = gen1.cuda()
gen2 = gen2.cuda()
criterion_GAN.cuda()
criterion_bce.cuda()
transform = transforms.Compose([
transforms.Resize((384, 384)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225))
def initialize(self, opt):
BaseModel.initialize(self, opt)
self.opt = opt
self.isTrain = opt.isTrain
# define tensors
self.input_A_1 = self.Tensor(opt.batchSize, opt.input_nc, opt.fineSize,
opt.fineSize)
self.input_A_2 = self.Tensor(opt.batchSize, opt.input_nc, opt.fineSize,
opt.fineSize)
self.input_B = self.Tensor(opt.batchSize, opt.output_nc, opt.fineSize,
opt.fineSize)
# load/define networks
self.netG = networks.define_G(opt.input_nc, opt.output_nc, opt.ngf,
opt.which_model_netG, opt.norm,
opt.use_dropout, self.gpu_ids)
self.netG.train()
if self.isTrain:
use_sigmoid = opt.no_lsgan
# D 的输入 是 9 channel了
if opt.which_model_netG == "sia_unet" or opt.which_model_netG == "stack_unet" or opt.which_model_netG == "sia_stack_unet":
self.netD = networks.define_D(
opt.input_nc + opt.input_nc + opt.output_nc, opt.ndf,
opt.which_model_netD, opt.n_layers_D, opt.norm,
use_sigmoid, self.gpu_ids)
else:
self.netD = networks.define_D(opt.input_nc + opt.output_nc,
opt.ndf, opt.which_model_netD,
opt.n_layers_D, opt.norm,
use_sigmoid, self.gpu_ids)
if not self.isTrain or opt.continue_train:
self.load_network(self.netG, 'G', opt.which_epoch)
if self.isTrain:
self.load_network(self.netD, 'D', opt.which_epoch)
if self.isTrain:
self.fake_AB_pool = ImagePool(opt.pool_size)
self.old_lr = opt.lr
# define loss functions
self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan,
tensor=self.Tensor)
self.criterionL1 = torch.nn.L1Loss()
# self.triLoss = networks.TRILoss(tensor=self.Tensor)
# need to implement
if self.opt.use_pose and self.opt.use_vgg:
# pose loss
self.criterionPOSE = networks.POSELoss()
# TVRegularizer loss
#self.criterionTVR = networks.TVRegularizerLoss()
# VGG loss
self.criterionVGG = VGGLoss()
# initialize optimizers
self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
self.optimizer_D = torch.optim.Adam(self.netD.parameters(),
lr=opt.lr,
betas=(opt.beta1, 0.999))
print('---------- Networks initialized ------------')
networks.print_network(self.netG)
networks.print_network(self.netD)
print('----------------------------------------------')
class Solver(object):
def __init__(self, data_loader, config):
"""Initialize configurations."""
self.data_loader = data_loader
self.config = config
self.build_model(config)
def build_model(self, config):
"""Create a generator and a discriminator."""
self.G = Generator(config.g_conv_dim, config.d_channel,
config.channel_1x1) # 2 for mask vector.
self.D = Discriminator(config.crop_size, config.d_conv_dim,
config.d_repeat_num)
self.G.apply(weights_init)
self.D.apply(weights_init)
self.G.cuda()
self.D.cuda()
self.g_optimizer = torch.optim.Adam(self.G.parameters(), config.g_lr,
[config.beta1, config.beta2])
self.d_optimizer = torch.optim.Adam(self.D.parameters(), config.d_lr,
[config.beta1, config.beta2])
self.G = nn.DataParallel(self.G)
self.D = nn.DataParallel(self.D)
self.VGGLoss = VGGLoss().eval()
self.VGGLoss.cuda()
self.VGGLoss = nn.DataParallel(self.VGGLoss)
def adv_loss(self, logits, target):
assert target in [1, 0]
targets = torch.full_like(logits, fill_value=target)
loss = F.binary_cross_entropy_with_logits(logits, targets)
return loss
def restore_model(self, resume_iters):
"""Restore the trained generator and discriminator."""
print(
'Loading the trained models from step {}...'.format(resume_iters))
G_path = os.path.join(self.model_save_dir,
'{}-G.ckpt'.format(resume_iters))
D_path = os.path.join(self.model_save_dir,
'{}-D.ckpt'.format(resume_iters))
self.G.load_state_dict(
torch.load(G_path, map_location=lambda storage, loc: storage))
self.D.load_state_dict(
torch.load(D_path, map_location=lambda storage, loc: storage))
def reset_grad(self):
"""Reset the gradient buffers."""
self.g_optimizer.zero_grad()
self.d_optimizer.zero_grad()
def update_lr(self, g_lr, d_lr):
"""Decay learning rates of the generator and discriminator."""
for param_group in self.g_optimizer.param_groups:
param_group['lr'] = g_lr
for param_group in self.d_optimizer.param_groups:
param_group['lr'] = d_lr
def denorm(self, x):
"""Convert the range from [-1, 1] to [0, 1]."""
out = (x + 1) / 2
return out.clamp_(0, 1)
def gradient_penalty(self, y, x):
"""Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
weight = torch.ones(y.size()).cuda()
dydx = torch.autograd.grad(outputs=y,
inputs=x,
grad_outputs=weight,
retain_graph=True,
create_graph=True,
only_inputs=True)[0]
dydx = dydx.view(dydx.size(0), -1)
dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
return torch.mean((dydx_l2norm - 1)**2)
def train(self):
data_loader = self.data_loader
config = self.config
# Learning rate cache for decaying.
g_lr = config.g_lr
d_lr = config.d_lr
# Label for lsgan
real_target = torch.full((config.batch_size, ), 1.).cuda()
fake_target = torch.full((config.batch_size, ), 0.).cuda()
criterion = nn.MSELoss().cuda()
# Start training.
print('Start training...')
start_time = time.time()
iteration = 0
num_iters_decay = config.num_epoch_decay * len(data_loader)
for epoch in range(config.num_epoch):
for i, (I_ori, I_gt, I_r, I_s) in enumerate(data_loader):
iteration += i
I_ori = I_ori.cuda(non_blocking=True)
I_gt = I_gt.cuda(non_blocking=True)
I_r = I_r.cuda(non_blocking=True)
I_s = I_s.cuda(non_blocking=True)
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# Compute loss with real images.
I_gt.requires_grad_(requires_grad=True)
out = self.D(I_gt)
# d_loss_real = criterion(out, real_target) * 0.5
d_loss_real = self.adv_loss(out, 1)
d_loss_reg = r1_reg(out, I_gt)
# Compute loss with fake images.
I_fake = self.G(I_r, I_s)
out = self.D(I_fake.detach())
# d_loss_fake = criterion(out, fake_target) * 0.5
d_loss_fake = self.adv_loss(out, 0)
# Backward and optimize.
d_loss = d_loss_real + d_loss_fake + d_loss_reg
self.reset_grad()
d_loss.backward()
self.d_optimizer.step()
# Logging.
loss = {}
loss['D/loss_real'] = d_loss_real.item()
loss['D/loss_fake'] = d_loss_fake.item()
loss['D/loss_reg'] = d_loss_reg.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
I_gt.requires_grad_(requires_grad=False)
# if (i+1) % config.n_critic == 0:
I_fake, g_loss_tr = self.G(I_r, I_s, IsGTrain=True)
out = self.D(I_fake)
# g_loss_fake = criterion(out, real_target)
g_loss_fake = self.adv_loss(out, 1)
g_loss_rec = torch.mean(torch.abs(I_fake - I_gt)) # Eq.(6)
g_loss_prec, g_loss_style = self.VGGLoss(I_gt, I_fake)
g_loss_prec *= config.lambda_perc
g_loss_style *= config.lambda_style
# Backward and optimize.
g_loss = g_loss_fake + config.lambda_rec * g_loss_rec + config.lambda_tr * g_loss_tr + g_loss_prec + g_loss_style
self.reset_grad()
g_loss.backward()
self.g_optimizer.step()
# Logging.
loss['G/loss_fake'] = g_loss_fake.item()
loss['G/loss_rec'] = g_loss_rec.item()
loss['G/loss_tr'] = g_loss_tr.item()
loss['G/loss_prec'] = g_loss_prec.item()
loss['G/loss_style'] = g_loss_style.item()
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
# Print out training information.
if (i + 1) % config.log_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
log = "Elapsed [{}], Epoch [{}/{}], Iteration [{}/{}], g_lr {:.5f}, d_lr {:.5f}".format(
et, epoch, config.num_epoch, i + 1, len(data_loader),
g_lr, d_lr)
for tag, value in loss.items():
log += ", {}: {:.4f}".format(tag, value)
print(log)
# Decay learning rates.
if (epoch + 1) > config.num_epoch_decay:
g_lr -= (config.g_lr / float(num_iters_decay))
d_lr -= (config.d_lr / float(num_iters_decay))
self.update_lr(g_lr, d_lr)
# print ('Decayed learning rates, g_lr: {}, d_lr: {}.'.format(g_lr, d_lr))
# Translate fixed images for debugging.
if (epoch + 1) % config.sample_epoch == 0:
with torch.no_grad():
I_fake_ori = self.G(I_ori, I_s)
I_fake_zero = self.G(torch.zeros(I_ori.size()), I_s)
sample_path = os.path.join(config.sample_dir,
'{}.jpg'.format(epoch))
I_concat = self.denorm(
torch.cat([
I_ori, I_gt, I_r, I_fake, I_fake_ori, I_fake_zero
],
dim=2))
I_concat = torch.cat(
[I_concat, I_s.repeat(1, 3, 1, 1)], dim=2)
save_image(I_concat.data.cpu(), sample_path)
print('Saved real and fake images into {}...'.format(
sample_path))
G_path = os.path.join(config.model_save_dir,
'{}-G.ckpt'.format(epoch + 1))
torch.save(self.G.state_dict(), G_path)
print('Saved model checkpoints into {}...'.format(
config.model_save_dir))
def test(self):
"""Translate images using StarGAN trained on a single dataset."""
# Load the trained generator.
self.restore_model(self.test_iters)
# Set data loader.
if self.dataset == 'CelebA':
data_loader = self.celeba_loader
elif self.dataset == 'RaFD':
data_loader = self.rafd_loader
with torch.no_grad():
for i, (x_real, c_org) in enumerate(data_loader):
# Prepare input images and target domain labels.
x_real = x_real.to(self.device)
c_trg_list = self.create_labels(c_org, self.c_dim,
self.dataset,
self.selected_attrs)
# Translate images.
x_fake_list = [x_real]
for c_trg in c_trg_list:
x_fake_list.append(self.G(x_real, c_trg))
# Save the translated images.
x_concat = torch.cat(x_fake_list, dim=3)
result_path = os.path.join(self.result_dir,
'{}-images.jpg'.format(i + 1))
save_image(self.denorm(x_concat.data.cpu()),
result_path,
nrow=1,
padding=0)
print('Saved real and fake images into {}...'.format(
result_path))