def main():
# Create the fundemental data loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[['/home/sunner/Music/waiting_for_you_dataset/wait'],
['/home/sunner/Music/waiting_for_you_dataset/real_world']],
transforms=transforms.Compose([
sunnertransforms.Resize((160, 320)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5]),
])),
batch_size=32,
shuffle=False,
num_workers=2)
# Use upper wrapper to assign particular iteration
loader = sunnerData.IterationLoader(loader, max_iter=1)
# Show!
for batch_tensor, _ in loader:
batch_img = sunnertransforms.asImg(batch_tensor, size=(160, 320))
cv2.imshow('show_window', batch_img[0][:, :, ::-1])
cv2.waitKey(0)
# Or show multiple image in one line
sunnertransforms.show(batch_tensor[:10], row=2, column=5)
def main():
# Define op first
transform_op = transforms.Compose([
sunnertransforms.Resize((160, 320)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
])
# Define loader
loader = sunnerData.DataLoader(sunnerData.VideoDataset(
root=[['/home/sunner/Music/single_flower_video_dataset/A'],
['/home/sunner/Music/single_flower_video_dataset/B']],
transforms=transform_op,
T=20),
batch_size=2,
shuffle=False,
num_workers=2)
# Looping
for _, seq in loader:
seq = [_.squeeze(1) for _ in torch.chunk(seq, seq.size(1), dim=1)
] # BTCHW -> T * BCHW
for i in range(10): # repeat for 10 times
for img in seq:
batch_img = sunnertransforms.asImg(img, size=(320, 640))
cv2.imshow('show_window', batch_img[0][:, :, ::-1])
cv2.waitKey(10)
break
def main():
# Split the dataset
train_dataset = sunnerData.ImageDataset(
root=[['/home/sunner/Music/waiting_for_you_dataset/wait'],
['/home/sunner/Music/waiting_for_you_dataset/real_world']],
transform=None,
split_ratio=0.1,
save_file=True)
del train_dataset
test_dataset = sunnerData.ImageDataset(
file_name='.split.pkl',
transform=transforms.Compose([
sunnertransforms.Resize((160, 320)),
sunnertransforms.ToTensor(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
]))
# Create the data loader
loader = sunnerData.DataLoader(test_dataset,
batch_size=32,
shuffle=False,
num_workers=2)
# Use upper wrapper to assign particular iteration
loader = sunnerData.IterationLoader(loader, max_iter=1)
# Show!
for batch_img, _ in loader:
batch_img = sunnertransforms.asImg(batch_img, size=(160, 320))
cv2.imshow('show_window', batch_img[0][:, :, ::-1])
cv2.waitKey(0)
def demo(args):
"""
This function define the demo process
Arg: args (napmespace) - The arguments
"""
# Create the data loader
loader = sunnerData.DataLoader(
dataset=sunnerData.ImageDataset(
root=[[args.demo]],
transforms=transforms.Compose([
sunnerTransforms.Resize(output_size=(args.H, args.W)),
sunnerTransforms.ToTensor(),
sunnerTransforms.ToFloat(),
# sunnerTransforms.Transpose(),
sunnerTransforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5]),
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
# Create the model
model = GANomaly2D(r=args.r, device=args.device)
model.IO(args.resume, direction='load')
# Demo!
bar = tqdm(loader)
model.eval()
with torch.no_grad():
for (img, ) in bar:
z, z_ = model.forward(img)
img, img_ = model.getImg()
visualizeAnomalyImage(img, img_, z, z_)
def main():
# Define the loader to generate the pallete object
loader = sunnerData.DataLoader(
sunnerData.ImageDataset(
root=[tag_folder],
transform=transforms.Compose([
sunnertransforms.ToTensor(),
]),
save_file=False # Don't save the record file, be careful!
),
batch_size=2,
shuffle=False,
num_workers=2)
pallete = sunnertransforms.getCategoricalMapping(loader,
path='pallete.json')[0]
del loader
# Define the actual loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[img_folder, tag_folder],
transform=transforms.Compose([
sunnertransforms.Resize((512, 1024)),
sunnertransforms.ToTensor(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=32,
shuffle=False,
num_workers=2)
# Define the reverse operator
goto_op = sunnertransforms.CategoricalTranspose(
pallete=pallete, direction=sunnertransforms.COLOR2ONEHOT)
back_op = sunnertransforms.CategoricalTranspose(
pallete=pallete, direction=sunnertransforms.ONEHOT2COLOR)
# Show!
for _, batch_index in loader:
batch_img = back_op(goto_op(batch_index))
batch_img = sunnertransforms.asImg(batch_img, size=(512, 1024))
cv2.imshow('show_window', batch_img[0][:, :, ::-1])
cv2.waitKey(0)
break
def main():
# Define the loader to generate the pallete object
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root = [
tag_folder
],
transforms = transforms.Compose([
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.UnNormalize(mean=[0, 0, 0], std=[255, 255, 255]), # Remember to transfer back to [0~255] before generate pallete
sunnertransforms.Transpose(sunnertransforms.BCHW2BHWC) # Remember to transfer back to BHWC before generate pallete
])
), batch_size = 2, shuffle = False, num_workers = 2
)
pallete = sunnertransforms.getCategoricalMapping(loader, path = 'pallete.json')[0]
del loader
# Define the actual loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root = [
img_folder,
tag_folder
],
transforms = transforms.Compose([
sunnertransforms.Resize((512, 1024)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Normalize(mean = [0.5, 0.5, 0.5], std = [0.5, 0.5, 0.5]),
])), batch_size = 32, shuffle = False, num_workers = 2
)
# Define the reverse operator
goto_op = sunnertransforms.CategoricalTranspose(pallete = pallete, direction = sunnertransforms.COLOR2ONEHOT)
back_op = sunnertransforms.CategoricalTranspose(pallete = pallete, direction = sunnertransforms.ONEHOT2COLOR)
# Show!
for _, batch_index in loader:
batch_img = back_op(goto_op(batch_index))
batch_img = sunnertransforms.asImg(batch_img, size = (512, 1024))
cv2.imshow('show_window', batch_img[0][:, :, ::-1])
cv2.waitKey(0)
break
def train(args):
"""
This function define the training process
Arg: args (napmespace) - The arguments
"""
# Create the data loader
loader = sunnerData.DataLoader(
dataset=sunnerData.ImageDataset(
root=[[args.train]],
transforms=transforms.Compose([
# transforms.RandomCrop(720,720)
# transforms.RandomRotation(45)
# transforms.RandomHorizontalFlip(),
# transforms.ColorJitter(brightness=0.5, contrast=0.5),
sunnerTransforms.Resize(output_size=(args.H, args.W)),
#transforms.RandomCrop(512,512)
sunnerTransforms.ToTensor(),
sunnerTransforms.ToFloat(),
# sunnerTransforms.Transpose(),
sunnerTransforms.Normalize(mean=[0.5, 0.5, 0.5],
std=[0.5, 0.5, 0.5]),
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
loader = sunnerData.IterationLoader(loader, max_iter=args.n_iter)
# Create the model
model = GANomaly2D(r=args.r, device=args.device)
model.IO(args.resume, direction='load')
model.train()
# Train!
bar = tqdm(loader)
for i, (normal_img, ) in enumerate(bar):
model.forward(normal_img)
model.backward()
loss_G, loss_D = model.getLoss()
bar.set_description("Loss_G: " + str(loss_G) + " loss_D: " +
str(loss_D))
bar.refresh()
if i % args.record_iter == 0:
model.eval()
with torch.no_grad():
z, z_ = model.forward(normal_img)
img, img_ = model.getImg()
visualizeEncoderDecoder(img, img_, z, z_, i)
model.train()
model.IO(args.det, direction='save')
model.IO(args.det, direction='save')
def main():
# Create the fundemental data loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[['/home/sunner/Music/waiting_for_you_dataset/wait'],
['/home/sunner/Music/waiting_for_you_dataset/real_world']],
transform=transforms.Compose([
sunnertransforms.Resize((160, 320)),
sunnertransforms.ToTensor(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=32,
shuffle=False,
num_workers=2)
# Use upper wrapper to assign particular iteration
loader = sunnerData.IterationLoader(loader, max_iter=1)
# Show!
for batch_img, _ in loader:
batch_img = sunnertransforms.asImg(batch_img, size=(160, 320))
cv2.imshow('show_window', batch_img[0][:, :, ::-1])
cv2.waitKey(0)
def main(opts):
# Create the data loader
loader = sunnerData.DataLoader(
sunnerData.ImageDataset(root=[[opts.path]],
transform=transforms.Compose([
sunnertransforms.Resize((1024, 1024)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(
sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=opts.batch_size,
shuffle=True,
)
# Create the model
start_epoch = 0
G = StyleGenerator()
D = StyleDiscriminator()
# Load the pre-trained weight
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
else:
INFO(
"Pre-trained weight cannot load successfully, train from scratch!")
# Multi-GPU support
if torch.cuda.device_count() > 1:
INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs")
G = nn.DataParallel(G)
D = nn.DataParallel(D)
G.to(opts.device)
D.to(opts.device)
# Create the criterion, optimizer and scheduler
optim_D = optim.Adam(D.parameters(), lr=0.00001, betas=(0.5, 0.999))
optim_G = optim.Adam(G.parameters(), lr=0.00001, betas=(0.5, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# Train
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
softplus = nn.Softplus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img, ) in enumerate(bar):
# =======================================================================================================
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
# =======================================================================================================
# Compute adversarial loss toward discriminator
D.zero_grad()
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device))
fake_logit = D(fake_img.detach())
d_loss = softplus(fake_logit).mean()
d_loss = d_loss + softplus(-real_logit).mean()
if opts.r1_gamma != 0.0:
r1_penalty = R1Penalty(real_img.detach(), D)
d_loss = d_loss + r1_penalty * (opts.r1_gamma * 0.5)
if opts.r2_gamma != 0.0:
r2_penalty = R2Penalty(fake_img.detach(), D)
d_loss = d_loss + r2_penalty * (opts.r2_gamma * 0.5)
loss_D_list.append(d_loss.item())
# Update discriminator
d_loss.backward()
optim_D.step()
# =======================================================================================================
# (2) Update G network: maximize log(D(G(z)))
# =======================================================================================================
if i % CRITIC_ITER == 0:
G.zero_grad()
fake_logit = D(fake_img)
g_loss = softplus(-fake_logit).mean()
loss_G_list.append(g_loss.item())
# Update generator
g_loss.backward()
optim_G.step()
# Output training stats
bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format(
ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
# Check how the generator is doing by saving G's output on fixed_noise
with torch.no_grad():
fake_img = G(fix_z).detach().cpu()
save_image(fake_img,
os.path.join(opts.det, 'images',
str(ep) + '.png'),
nrow=4,
normalize=True)
# Save model
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
def train(opts):
def log(string, name="stylegan.log"):
with open(name, 'a') as f:
f.write(string + '\n')
writer = SummaryWriter(str(opts.output))
loader = dataset.DataLoader(dataset=dataset.ImageDataset(
[[opts.input]],
transform=transforms.Compose([
trans.Resize((opts.imsize, opts.imsize)),
trans.ToTensor(),
trans.ToFloat(),
trans.Transpose(trans.BHWC2BCHW),
trans.Normalize()
])),
batch_size=opts.batch_size,
shuffle=True)
G = opts.G
D = opts.D
step = 0
start_epoch = opts.start_epoch
if opts.resume:
try:
assert os.path.exists(opts.resume)
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
logger.info("Load Pretrained Weight")
except:
logger.warn("Resume Files cannot Load")
logger.info("Train from Scratch")
else:
logger.info("Train from Scratch")
if torch.cuda.device_count() > 1 and opts.device == 'cuda':
logger.info(f"{torch.cuda.device_count()} GPUs found.")
G = nn.DataParallel(G)
D = nn.DataParallel(D)
if opts.device == 'cuda':
torch.backends.cudnn.benchmark = True
G.to(opts.device)
D.to(opts.device)
optimG = Adam(G.parameters(), lr=opts.g_lr, betas=opts.betas)
optimD = Adam(D.parameters(), lr=opts.d_lr, betas=opts.betas)
schedulerG = lr_scheduler.ExponentialLR(optimG, gamma=opts.g_lrdecay)
schedulerD = lr_scheduler.ExponentialLR(optimD, gamma=opts.d_lrdecay)
fixed_z = torch.randn([opts.batch_size, 512]).to(opts.device)
sp = nn.Softplus()
g_store = [0.0]
d_store = [0.0]
for epoch in range(start_epoch, opts.epochs + 1):
bar = tqdm(loader)
glosses = []
dlosses = []
for i, (real, ) in enumerate(bar):
step += 1
D.zero_grad()
real = real.to(opts.device)
Dr = D(real)
writer.add_graph(D, real)
z = torch.randn([real.size(0), 512]).to(opts.device)
fake = G(z)
writer.add_graph(G, z)
Df = D(fake.detach())
Dloss = sp(Df).mean() + sp(-Dr).mean()
if opts.r1gamma > 0:
r1 = r1_penalty(real.detach(), D)
Dloss = Dloss + r1 * (opts.r1gamma * .5)
if opts.r2gamma > 0:
r2 = r2_penalty(fake.detach(), D)
Dloss = Dloss + r2 * (opts.r2gamma * .5)
dlosses.append(Dloss.item())
Dloss.backward()
optimD.step()
if i % opts.critic_iters == 0:
G.zero_grad()
Df = D(fake)
Gloss = sp(-Df).mean()
glosses.append(Gloss.item())
Gloss.backward()
optimG.step()
if i % opts.show_interval == 0:
with torch.no_grad():
nr = int(math.ceil(math.sqrt(opts.batch_size)))
z = torch.randn([real.size(0), 512]).to(opts.device)
img = G(z)
save_image(img.detach().cpu(),
os.path.join(opts.output, 'images', 'normal',
f'{epoch:04}_{i:06}.png'),
nrow=nr,
normalize=True)
fakes = utils.make_grid(img, nr, padding=0)
fakes = fakes.to(torch.float32).cpu().numpy()
fakes = np.clip((fakes / 2) + 0.5, 0, 1)
writer.add_image(f"EPOCH{epoch}/Random",
torch.from_numpy(fakes), i)
img = G(fixed_z)
save_image(img.detach().cpu(),
os.path.join(opts.output, 'images', 'fixed',
f'{epoch:04}_{i:06}.png'),
nrow=nr,
normalize=True)
fakes = utils.make_grid(img, nr, padding=0)
fakes = fakes.to(torch.float32).cpu().numpy()
fakes = np.clip((fakes / 2) + 0.5, 0, 1)
writer.add_image(f"EPOCH{epoch}/Fixed",
torch.from_numpy(fakes), i)
writer.add_scalar(f"LOSS/Generator",
Gloss.item(),
global_step=step)
writer.add_scalar(f"LOSS/Discriminator",
Dloss.item(),
global_step=step)
bar.set_description(
f"Epoch {epoch}/{opts.epochs} G: {glosses[-1]:.6f} D: {dlosses[-1]:.6f}"
)
g_store.append(np.mean(glosses))
d_store.append(np.mean(dlosses))
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': g_store,
'Loss_D': d_store,
'start_epoch': epoch,
'opts': opts
}
torch.save(state, os.path.join(opts.output, 'models', 'latest.pth'))
if epoch % 10 == 0:
torch.save(state,
os.path.join(opts.output, 'models', f'{epoch:04}.pth'))
schedulerD.step()
schedulerG.step()
def main(opts):
# Create the data loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[[opts.path]],
transform=transforms.Compose([
sunnertransforms.Resize((128, 128)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=opts.batch_size,
shuffle=True,
num_workers=4)
# Create the model
if opts.type == 'style':
G = StyleGenerator().to(opts.device)
else:
G = Generator().to(opts.device)
D = Discriminator().to(opts.device)
# Load the pre-trained weight
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
else:
INFO(
"Pre-trained weight cannot load successfully, train from scratch!")
# Create the criterion, optimizer and scheduler
optim_D = optim.Adam(D.parameters(), lr=0.0001, betas=(0.5, 0.999))
optim_G = optim.Adam(G.parameters(), lr=0.0001, betas=(0.5, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# Train
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img, ) in enumerate(bar):
# =======================================================================================================
# Update discriminator
# =======================================================================================================
# Compute adversarial loss toward discriminator
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img = G(torch.randn([real_img.size(0), 512]).to(opts.device))
fake_logit = D(fake_img.detach())
d_loss = -(real_logit.mean() -
fake_logit.mean()) + gradient_penalty(
real_img.data, fake_img.data, D) * 10.0
loss_D_list.append(d_loss.item())
# Update discriminator
optim_D.zero_grad()
d_loss.backward()
optim_D.step()
# =======================================================================================================
# Update generator
# =======================================================================================================
if i % CRITIC_ITER == 0:
# Compute adversarial loss toward generator
fake_img = G(
torch.randn([opts.batch_size, 512]).to(opts.device))
fake_logit = D(fake_img)
g_loss = -fake_logit.mean()
loss_G_list.append(g_loss.item())
# Update generator
D.zero_grad()
optim_G.zero_grad()
g_loss.backward()
optim_G.step()
bar.set_description(" {} [G]: {} [D]: {}".format(
ep, loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
fake_img = G(fix_z)
save_image(fake_img,
os.path.join(opts.det, 'images',
str(ep) + '.png'),
nrow=4,
normalize=True)
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
def main(opts):
# Create the data loader
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[[opts.path]],
transform=transforms.Compose([
sunnertransforms.Resize((opts.resolution, opts.resolution)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=opts.batch_size,
shuffle=True,
drop_last=True
)
# Create the model
start_epoch = 0
G = G_stylegan2(fmap_base=opts.fmap_base,
resolution=opts.resolution,
mapping_layers=opts.mapping_layers,
opts=opts,
return_dlatents=True)
D = D_stylegan2(fmap_base=opts.fmap_base,
resolution=opts.resolution,
structure='resnet')
# Load the pre-trained weight
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
else:
INFO("Pre-trained weight cannot load successfully, train from scratch!")
# Multi-GPU support
if torch.cuda.device_count() > 1:
INFO("Multiple GPU:" + str(torch.cuda.device_count()) + "\t GPUs")
G = torch.nn.DataParallel(G)
D = torch.nn.DataParallel(D)
G.to(opts.device)
D.to(opts.device)
# Create the criterion, optimizer and scheduler
lr_D = 0.0015
lr_G = 0.0015
optim_D = torch.optim.Adam(D.parameters(), lr=lr_D, betas=(0.9, 0.999))
# g_mapping has 100x lower learning rate
params_G = [{"params": G.g_synthesis.parameters()},
{"params": G.g_mapping.parameters(), "lr": lr_G * 0.01}]
optim_G = torch.optim.Adam(params_G, lr=lr_G, betas=(0.9, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# Train
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
softplus = torch.nn.Softplus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img,) in enumerate(bar):
real_img = real_img.to(opts.device)
latents = torch.randn([real_img.size(0), 512]).to(opts.device)
# =======================================================================================================
# (1) Update D network: D_logistic_r1(default)
# =======================================================================================================
# Compute adversarial loss toward discriminator
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img, fake_dlatent = G(latents)
fake_logit = D(fake_img.detach())
d_loss = softplus(fake_logit)
d_loss = d_loss + softplus(-real_logit)
# original
r1_penalty = D_logistic_r1(real_img.detach(), D)
d_loss = (d_loss + r1_penalty).mean()
# lite
# d_loss = d_loss.mean()
loss_D_list.append(d_loss.mean().item())
# Update discriminator
optim_D.zero_grad()
d_loss.backward()
optim_D.step()
# =======================================================================================================
# (2) Update G network: G_logistic_ns_pathreg(default)
# =======================================================================================================
# if i % CRITIC_ITER == 0:
G.zero_grad()
fake_scores_out = D(fake_img)
_g_loss = softplus(-fake_scores_out)
# Compute |J*y|.
# pl_noise = (torch.randn(fake_img.shape) / np.sqrt(fake_img.shape[2] * fake_img.shape[3])).to(fake_img.device)
# pl_grads = grad(torch.sum(fake_img * pl_noise), fake_dlatent, retain_graph=True)[0]
# pl_lengths = torch.sqrt(torch.sum(torch.sum(torch.mul(pl_grads, pl_grads), dim=2), dim=1))
# pl_mean = PL_DECAY * torch.sum(pl_lengths)
#
# pl_penalty = torch.mul(pl_lengths - pl_mean, pl_lengths - pl_mean)
# reg = pl_penalty * PL_WEIGHT
#
# # original
# g_loss = (_g_loss + reg).mean()
# lite
g_loss = _g_loss.mean()
loss_G_list.append(g_loss.mean().item())
# Update generator
g_loss.backward(retain_graph=True)
optim_G.step()
# Output training stats
bar.set_description(
"Epoch {} [{}, {}] [G]: {} [D]: {}".format(ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1]))
# Save the result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
# Check how the generator is doing by saving G's output on fixed_noise
with torch.no_grad():
fake_img = G(fix_z)[0].detach().cpu()
save_image(fake_img, os.path.join(opts.det, 'images', str(ep) + '.png'), nrow=4, normalize=True)
# Save model
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
# Plot the total loss curve
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
def main(opts):
# Data load
loader = sunnerData.DataLoader(sunnerData.ImageDataset(
root=[[opts.path]],
transform=transforms.Compose([
sunnertransforms.Resize((opts.resolution, opts.resolution)),
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize()
])),
batch_size=opts.batch_size,
shuffle=True,
drop_last=True)
# model generation
start_epoch = 0
G = Generator_stylegan2(fmap_base=opts.fmap_base,
resol=opts.resolution,
mapping_layers=opts.mapping_layers,
opts=opts,
return_dlatents=True)
D = Discriminator_stylegan2(fmap_base=opts.fmap_base,
resol=opts.resolution,
structure='resnet')
# pre-trained weight loading
if os.path.exists(opts.resume):
INFO("Load the pre-trained weight!")
state = torch.load(opts.resume)
G.load_state_dict(state['G'])
D.load_state_dict(state['D'])
start_epoch = state['start_epoch']
else:
INFO("pre-trained weight error")
# multiple GPU support
if (torch.cuda.device_count() > 1):
INFO("multiple GPU detected! Total " + str(torch.cuda.device_count()) +
'\t GPUs!')
G = torch.nn.DataParrlel(G)
D = torch.nn.DataParallel(D)
G.to(opts.device)
D.to(opts.device)
# optimizer, scheduler
lr_D = 0.0015
lr_G = 0.0015
optim_D = torch.optim.Adam(D.parameters(), lr=lr_D, betas=(0.9, 0.999))
params_G = [{
"params": G.g_synthesis.parameters()
}, {
"params": G.g_mapping.parameters(),
"lr": lr_G * 0.01
}]
optim_G = torch.optim.Adam(params_G, lr=lr_G, betas=(0.9, 0.999))
scheduler_D = optim.lr_scheduler.ExponentialLR(optim_D, gamma=0.99)
scheduler_G = optim.lr_scheduler.ExponentialLR(optim_G, gamma=0.99)
# start training
fix_z = torch.randn([opts.batch_size, 512]).to(opts.device)
softplus = torch.nn.Softplus()
Loss_D_list = [0.0]
Loss_G_list = [0.0]
for ep in range(start_epoch, opts.epoch):
bar = tqdm(loader)
loss_D_list = []
loss_G_list = []
for i, (real_img, ) in enumerate(bar):
real_img = real_img.to(opts.device)
latents = torch.randn([real_img.size(0), 512]).to(opts.device)
# Discriminator Network
real_img = real_img.to(opts.device)
real_logit = D(real_img)
fake_img, fake_dlatent = G(latents)
fake_logit = D(fake_img.detach())
d_loss = softplus(fake_logit)
d_loss = d_loss + softplus(-real_logit)
r1_penalty = D_logistic_r1(real_img.detach(), D)
d_loss = (d_loss + r1_penalty).mean()
loss_D_list.append(d_loss.mean().item())
optim_D.zero_grad()
d_loss.backward()
optim_D.step()
# Generator Network
G.zero_grad()
fake_scores_out = D(fake_img)
_g_loss = softplus(-fake_scores_out)
g_loss = _g_loss.mean()
loss_G_list.append(g_loss.mean().item())
g_loss.backward()
optim_G.step()
bar.set_description("Epoch {} [{}, {}] [G]: {} [D]: {}".format(
ep, i + 1, len(loader), loss_G_list[-1], loss_D_list[-1]))
# save result
Loss_G_list.append(np.mean(loss_G_list))
Loss_D_list.append(np.mean(loss_D_list))
with torch.no_grad():
fake_img = G(fix_z)[0].detach().cpu()
save_image(fake_img,
os.path.join(opts.det, 'images',
str(ep) + '.png'),
nrow=4,
normalize=True)
# save model
state = {
'G': G.state_dict(),
'D': D.state_dict(),
'Loss_G': Loss_G_list,
'Loss_D': Loss_D_list,
'start_epoch': ep,
}
torch.save(state, os.path.join(opts.det, 'models', 'latest.pth'))
scheduler_D.step()
scheduler_G.step()
Loss_D_list = Loss_D_list[1:]
Loss_G_list = Loss_G_list[1:]
plotLossCurve(opts, Loss_D_list, Loss_G_list)
def train(opts):
# Create the loader
loader = sunnerData.DataLoader(dataset=BracketedDataset(
root=opts.folder,
crop_size=opts.crop_size,
transform=transforms.Compose([
sunnertransforms.ToTensor(),
sunnertransforms.ToFloat(),
sunnertransforms.Transpose(sunnertransforms.BHWC2BCHW),
sunnertransforms.Normalize(),
])),
batch_size=opts.batch_size,
shuffle=True,
num_workers=0)
# Create the model
model = DeepFuse(device=opts.device)
criterion = MEF_SSIM_Loss().to(opts.device)
optimizer = Adam(model.parameters(), lr=0.0001)
# Load pre-train model
if os.path.exists(opts.resume):
state = torch.load(opts.resume)
Loss_list = state['loss']
model.load_state_dict(state['model'])
else:
Loss_list = []
# Train
bar = tqdm(range(opts.epoch))
for ep in bar:
loss_list = []
for (patch1, patch2) in loader:
# Extract the luminance and move to computation device
patch1, patch2 = patch1.to(opts.device), patch2.to(opts.device)
patch1_lum = patch1[:, 0:1]
patch2_lum = patch2[:, 0:1]
# Forward and compute loss
model.setInput(patch1_lum, patch2_lum)
y_f = model.forward()
loss, y_hat = criterion(y_1=patch1_lum, y_2=patch2_lum, y_f=y_f)
loss_list.append(loss.item())
bar.set_description("Epoch: %d Loss: %.6f" % (ep, loss_list[-1]))
# Update the parameters
optimizer.zero_grad()
loss.backward()
optimizer.step()
Loss_list.append(np.mean(loss_list))
# Save the training image
if ep % opts.record_epoch == 0:
img = fusePostProcess(y_f, y_hat, patch1, patch2, single=False)
cv2.imwrite(os.path.join(opts.det, 'image',
str(ep) + ".png"), img[0, :, :, :])
# Save the training model
if ep % (opts.epoch // 5) == 0:
model_name = str(ep) + ".pth"
else:
model_name = "latest.pth"
state = {'model': model.state_dict(), 'loss': Loss_list}
torch.save(state, os.path.join(opts.det, 'model', model_name))
# Plot the loss curve
plt.clf()
plt.plot(Loss_list, '-')
plt.title("loss curve")
plt.savefig(os.path.join(opts.det, 'image', "curve.png"))
def train(args, dataset, generator, discriminator):
step = int(math.log2(args.init_size)) - 2
resolution = 4 * 2**step
data_loader = sunnerData.DataLoader(dataset,
batch_size=args.batch.get(
resolution, args.batch_default),
resolution=resolution,
shuffle=True)
adjust_lr(g_optimizer, args.lr.get(resolution, 0.001))
adjust_lr(d_optimizer, args.lr.get(resolution, 0.001))
pbar = tqdm(data_loader)
requires_grad(generator, False)
requires_grad(discriminator, True)
disc_loss_val = 0
gen_loss_val = 0
grad_loss_val = 0
alpha = 0
used_sample = 0
max_step = int(math.log2(args.max_size)) - 2
final_progress = False
for i, (real_img, ) in enumerate(bar):
discriminator.zero_grad()
alpha = min(1, 1 / args.phase * (used_sample + 1))
if (resolution == args.init_size
and args.ckpt is None) or final_progress:
alpha = 1
if used_sample > args.phase * 2:
used_sample = 0
step += 1
if step > max_step:
step = max_step
final_progress = True
ckpt_step = step + 1
else:
alpha = 0
ckpt_step = step
resolution = 4 * 2**step
data_loader = sunnerData.DataLoader(dataset,
batch_size=args.batch.get(
resolution,
args.batch_default),
resolution=resolution,
shuffle=True)
torch.save(
{
"generator": generator.module.state_dict(),
"discriminator": discriminator.module.state_dict(),
"g_optimizer": g_optimizer.state_dict(),
"d_optimizer": d_optimizer.state_dict(),
"g_running": g_running.state_dict(),
},
f"checkpoint/train_step-{ckpt_step}.model",
)
adjust_lr(g_optimizer, args.lr.get(resolution, 0.001))
adjust_lr(d_optimizer, args.lr.get(resolution, 0.001))
try:
real_image = next(data_loader)
except (OSError, StopIteration):
data_loader = iter(loader)
real_image = next(data_loader)
used_sample += real_image.shape[0]
b_size = real_image.size(0)
real_image = real_image.cuda()
real_predict = discriminator(real_image, step=step, alpha=alpha)
real_predict = real_predict.mean() - 0.001 * (real_predict**2).mean()
(-real_predict).backward()
if args.mixing and random.random() < 0.9:
gen_in11, gen_in12, gen_in21, gen_in22 = torch.randn(
4, b_size, code_size, device="cuda").chunk(4, 0)
gen_in1 = [gen_in11.squeeze(0), gen_in12.squeeze(0)]
gen_in2 = [gen_in21.squeeze(0), gen_in22.squeeze(0)]
else:
gen_in1, gen_in2 = torch.randn(2, b_size, code_size,
device="cuda").chunk(2, 0)
gen_in1 = gen_in1.squeeze(0)
gen_in2 = gen_in2.squeeze(0)
fake_image = generator(gen_in1, step=step, alpha=alpha)
fake_predict = discriminator(fake_image, step=step, alpha=alpha)
fake_predict = fake_predict.mean()
fake_predict.backward()
eps = torch.rand(b_size, 1, 1, 1).cuda()
x_hat = eps * real_image.data + (1 - eps) * fake_image.data
x_hat.requires_grad = True
hat_predict = discriminator(x_hat, step=step, alpha=alpha)
grad_x_hat = grad(outputs=hat_predict.sum(),
inputs=x_hat,
create_graph=True)[0]
grad_penalty = (
(grad_x_hat.view(grad_x_hat.size(0), -1).norm(2, dim=1) -
1)**2).mean()
grad_penalty = 10 * grad_penalty
grad_penalty.backward()
if i % 10 == 0:
grad_loss_val = grad_penalty.item()
disc_loss_val = (real_predict - fake_predict).item()
d_optimizer.step()
if (i + 1) % n_critic == 0:
generator.zero_grad()
requires_grad(generator, True)
requires_grad(discriminator, False)
fake_image = generator(gen_in2, step=step, alpha=alpha)
predict = discriminator(fake_image, step=step, alpha=alpha)
if args.loss == "wgan-gp":
loss = -predict.mean()
elif args.loss == "r1":
loss = F.softplus(-predict).mean()
if i % 10 == 0:
gen_loss_val = loss.item()
loss.backward()
g_optimizer.step()
accumulate(g_running, generator.module)
requires_grad(generator, False)
requires_grad(discriminator, True)
if (i + 1) % 10000 == 0:
images = []
gen_i, gen_j = args.gen_sample.get(resolution, (10, 5))
with torch.no_grad():
for _ in range(gen_i):
images.append(
g_running(torch.randn(gen_j, code_size).cuda(),
step=step,
alpha=alpha).data.cpu())
utils.save_image(torch.cat(images, 0),
f"sample/{str(i + 1).zfill(6)}.png",
nrow=gen_i,
normalize=True,
range=(-1, 1))
torch.save(g_running.state_dict(),
f"checkpoint/{str(i + 1).zfill(6)}.model")
state_msg = f"Size: {4 * 2 ** step}; G: {gen_loss_val:.3f}; D: {disc_loss_val:.3f};" f" Grad: {grad_loss_val:.3f}; Alpha: {alpha:.5f}"
pbar.set_description(state_msg)