def test_pix2pix(pathdataset: str, pathoutput: str, pathmodel: str) -> str:
"""
Use the Pix2Pix trained model to generate ("translate")
images from the Test dataset
Args:
pathdataset (string) : Full path name to the test dataset (TFRecords file)
pathoutput (string) : Full path name to the Output ("translated") images
pathmodel (string) : Full path name to the Pix2Pix model checkpoint
Returns:
output_path (string) : Path to generated/translated images
"""
# ------------------------------
# In order to be able to convert
# a Python Function directly
# into a Kubeflow component,
# we need to move the python
# includes inside that python
# function.
# ------------------------------
import numpy as np
import tensorflow as tf
# The below tag comment is used by a tool script from this project to automatically nest
# the python code of the imports function tagged KFP-NEST-IMPORT, just right after this tag.
# This is only usefull when using the SDK's kfp.components.func_to_container_op method,
# which allows to convert a Python function to a pipeline component and returns a factory function
#
#KFP-NEST-HERE
# ------------------------------
# Define some hyperparameters
# (Not managed as Kubeflow
# pipeline parameters)
# ------------------------------
BATCH_SIZE = 1
# ------------------------------
# Restore the Pix2Pix GAN
#
# Rebuild the Generator
# (Skip the Discriminator
# and Optimizers)
# ------------------------------
generator = Generator()
# Layer objects in TensorFlow may delay the creation of variables
# to their first call, when input shapes are available.
_ = generator(
np.random.uniform(-127.5, 127.5, (1, 256, 256, 3)).astype(np.float32))
# ------------------------------
# Restore the last Checkpoint
# (Skip the Discriminator
# and Optimizers)
# -----------------------------
# Configure the model checkpoints saving
checkpoint = tf.train.Checkpoint(generator=generator)
# Restoring the latest checkpoint
manager = tf.train.CheckpointManager(checkpoint,
directory=pathmodel,
max_to_keep=1,
checkpoint_name='ckpt')
_ = checkpoint.restore(manager.latest_checkpoint).expect_partial()
if manager.latest_checkpoint:
print("Model restored from {}".format(manager.latest_checkpoint))
else:
print("[WARNING] Initializing model from scratch.")
# ------------------------------
# One step Inference
# nested Helper function
# ------------------------------
#@tf.function # Compile function into a graph for faster execution
# TODO : Using Autograd does not seem to work with Checkpoint Restore ??
def generate_image(img):
""" Generate an image using the Pix2pix Generator """
return generator(img, training=False)
# ------------------------------
# Extract the Test Dataset
# from the TFRecords file
# ------------------------------
test_dataset = get_dataset(pathdataset, BATCH_SIZE, shuffle=False)
# ------------------------------
# Loop over the training batches
# ------------------------------
for image_features in test_dataset:
# Loop over all images in the batch
for record_idx in range(BATCH_SIZE):
# ------------------------------
# Extract the individual
# features and reconstruct
# the input images to feed
# the Neural Networks
# ------------------------------
a_image, b_image, file_name = decode_tfrecords_example(
image_features, record_idx)
print("Processing image ", file_name)
# ------------------------------
# Apply same Data preparation,
# as for the training, on
# b_image - the input image -
# (but without Data augmentation)
# in order to feed into the
# Generator Neural Network Input
# ------------------------------
img_b = normalize(b_image) #[ height, width, OUTPUT_CHANNELS]
img_b = tf.image.resize(img_b, [256, 256])
# Add a batch dimension to have 4D Tensor images to feed the Neural Network input
img_b = tf.expand_dims(img_b,
axis=0) # [ 1, height, width, channel]
# ------------------------------
# Generate output image
# - One step inference -
# ------------------------------
fake_b = generate_image(img_b)
# ------------------------------
# Save the source, target
# and generated images to disk
# (incl. decoding if necessary)
# ------------------------------
img_name = pathoutput + "/img_a-" + file_name + ".jpg"
save_image(a_image, img_name)
img_name = pathoutput + "/img_b-" + file_name + ".jpg"
save_image(b_image, img_name)
img_name = pathoutput + "/fake_b-" + file_name + ".jpg"
save_image(fake_b, img_name)
print("End of inference on Test Dataset")
# ------------------------------
# Write the Output of the
# Kubeflow Pipeline Component
# ------------------------------
try:
# This works only inside Docker containers
with open('/output.txt', 'w') as f:
f.write(pathoutput)
except PermissionError:
pass
return pathoutput
def project(path_ckpt, path_files, step=1000):
device = "cuda"
parser = argparse.ArgumentParser()
parser.add_argument('-f', type=str, help='jup kernel')
# parser.add_argument("--ckpt", type=str, required=True)
parser.add_argument("--size", type=int, default=512)
parser.add_argument("--lr_rampup", type=float, default=0.05)
parser.add_argument("--lr_rampdown", type=float, default=0.25)
parser.add_argument("--lr", type=float, default=0.1)
parser.add_argument("--noise", type=float, default=0.05)
parser.add_argument("--noise_ramp", type=float, default=0.75)
parser.add_argument("--step", type=int, default=1000)
parser.add_argument("--noise_regularize", type=float, default=1e5)
parser.add_argument("--mse", type=float, default=0)
# parser.add_argument("--w_plus", action="store_true")
# parser.add_argument("files", metavar="FILES", nargs="+")
args = parser.parse_args()
args.ckpt = path_ckpt
args.files = path_files
args.w_plus = False
args.step = step
n_mean_latent = 10000
resize = min(args.size, 256)
transform = transforms.Compose([
transforms.Resize(resize),
transforms.CenterCrop(resize),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])
imgs = []
for imgfile in args.files:
img = transform(Image.open(imgfile).convert("RGB"))
imgs.append(img)
imgs = torch.stack(imgs, 0).to(device)
g_ema = Generator(args.size, 512, 8)
g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False)
g_ema.eval()
g_ema = g_ema.to(device)
with torch.no_grad():
noise_sample = torch.randn(n_mean_latent, 512, device=device)
latent_out = g_ema.style(noise_sample)
latent_mean = latent_out.mean(0)
latent_std = ((latent_out - latent_mean).pow(2).sum() / n_mean_latent) ** 0.5
percept = lpips.PerceptualLoss(
model="net-lin", net="vgg", use_gpu=device.startswith("cuda")
)
noises_single = g_ema.make_noise()
noises = []
for noise in noises_single:
noises.append(noise.repeat(imgs.shape[0], 1, 1, 1).normal_())
latent_in = latent_mean.detach().clone().unsqueeze(0).repeat(imgs.shape[0], 1)
if args.w_plus:
latent_in = latent_in.unsqueeze(1).repeat(1, g_ema.n_latent, 1)
latent_in.requires_grad = True
for noise in noises:
noise.requires_grad = True
optimizer = optim.Adam([latent_in] + noises, lr=args.lr)
pbar = tqdm(range(args.step))
latent_path = []
for i in pbar:
t = i / args.step
lr = get_lr(t, args.lr)
optimizer.param_groups[0]["lr"] = lr
noise_strength = latent_std * args.noise * max(0, 1 - t / args.noise_ramp) ** 2
latent_n = latent_noise(latent_in, noise_strength.item())
img_gen, _ = g_ema([latent_n], input_is='latent', noise=noises)
batch, channel, height, width = img_gen.shape
if height > resize:
factor = height // resize
img_gen = img_gen.reshape(
batch, channel, height // factor, factor, width // factor, factor
)
img_gen = img_gen.mean([3, 5])
p_loss = percept(img_gen, imgs).sum()
n_loss = noise_regularize(noises)
mse_loss = F.mse_loss(img_gen, imgs)
loss = p_loss + args.noise_regularize * n_loss + args.mse * mse_loss
optimizer.zero_grad()
loss.backward()
optimizer.step()
noise_normalize_(noises)
if (i + 1) % 100 == 0:
latent_path.append(latent_in.detach().clone())
pbar.set_description((
f"perceptual: {p_loss.item():.8f}; noise regularize: {n_loss.item():.8f}; mse: {mse_loss.item():.8f}; lr: {lr:.4f}"
))
img_gen, _ = g_ema([latent_path[-1]], input_is='latent', noise=noises)
filename = os.path.splitext(os.path.basename(args.files[0]))[0] + ".pt"
img_ar = make_image(img_gen)
result_file = {}
for i, input_name in enumerate(args.files):
noise_single = []
for noise in noises:
noise_single.append(noise[i: i + 1])
result_file[input_name] = {
"img": img_gen[i],
"latent": latent_in[i],
"noise": noise_single,
}
img_name = os.path.splitext(os.path.basename(input_name))[0] + "-project.png"
pil_img = Image.fromarray(img_ar[i])
pil_img.save(img_name)
torch.save(result_file, filename)
print(filename)
return img_gen, latent_path, latent_in
def train():
data, info = tfds.load("mnist",
with_info=True,
data_dir='/data/tensorflow_datasets')
train_data = data['train']
if not os.path.exists('./images'):
os.makedirs('./images')
# settting hyperparameter
latent_dim = 100
epochs = 800
batch_size = 200
buffer_size = 6000
save_interval = 50
generator = Generator()
discriminator = Discriminator()
gen_optimizer = tf.keras.optimizers.Adam(0.0002)
disc_optimizer = tf.keras.optimizers.Adam(0.0002)
train_dataset = train_data.map(normalize).shuffle(buffer_size).batch(
batch_size)
cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)
@tf.function
def train_step(images):
noise = tf.random.normal([batch_size, latent_dim])
with tf.GradientTape(persistent=True) as tape:
generated_images = generator(noise)
real_output = discriminator(images)
generated_output = discriminator(generated_images)
gen_loss = generator_loss(cross_entropy, generated_output)
disc_loss = discriminator_loss(cross_entropy, real_output,
generated_output)
grad_disc = tape.gradient(disc_loss, discriminator.trainable_variables)
grad_gen = tape.gradient(gen_loss, generator.trainable_variables)
disc_optimizer.apply_gradients(
zip(grad_disc, discriminator.trainable_variables))
gen_optimizer.apply_gradients(
zip(grad_gen, generator.trainable_variables))
return gen_loss, disc_loss
seed = tf.random.normal([16, latent_dim])
for epoch in range(epochs):
start = time.time()
total_gen_loss = 0
total_disc_loss = 0
for images in train_dataset:
gen_loss, disc_loss = train_step(images)
total_gen_loss += gen_loss
total_disc_loss += disc_loss
print('Time for epoch {} is {} sec - gen_loss = {}, disc_loss = {}'.
format(epoch + 1,
time.time() - start, total_gen_loss / batch_size,
total_disc_loss / batch_size))
if epoch % save_interval == 0:
save_imgs(epoch, generator, seed)
import torch
from torch import optim
from torch.autograd import Variable
import torchvision
import os
import matplotlib.pyplot as plt
from model import Generator, Discriminator
from utils import same_seeds
from data import get_dataset
from torch.utils.data import Dataset, DataLoader
import torch.nn as nn
import sys
if __name__ == "__main__":
# hyperparameters
batch_size = 64
z_dim = 100
lr = 1e-4
n_epoch = 30
model_name = sys.argv[1]
file_name = sys.argv[2]
same_seeds(0)
# model
G = Generator(in_dim=z_dim).cuda()
G.load_state_dict(torch.load(model_name))
G.eval()
z_sample = Variable(torch.randn(100, z_dim)).cuda()
f_imgs_sample = (G(z_sample).data + 1) / 2.0
torchvision.utils.save_image(f_imgs_sample, file_name, nrow=10)
from model import Generator, Discriminator
MAX_EPOCHS = 500
LAMBDA = 1
BATCH_SIZE = 4
transform = transforms.Compose([ToTensor()])
train_set = FacadeDataset("datasets/facades/train", transform=transform)
val_set = FacadeDataset("datasets/facades/val", transform=transform)
test_set = FacadeDataset("datasets/facades/test", transform=transform)
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, \
shuffle=True, num_workers=0)
generator = Generator().cuda()
discriminator = Discriminator().cuda()
gen_optim = torch.optim.Adam(generator.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
dis_optim = torch.optim.Adam(discriminator.parameters(),
lr=0.0002,
betas=(0.5, 0.999))
gen_loss = torch.tensor(0)
dis_loss = torch.tensor(0)
for epoch in range(MAX_EPOCHS):
pbar = tqdm(train_loader, dynamic_ncols=True)
pbar.set_description(f'e{epoch}')
for idx, batch in enumerate(pbar):
# test_path = '../dataset/crop_people/test.txt'
show_result = './prediction'
if not os.path.exists(show_result):
os.mkdir(show_result)
# dirname = os.path.dirname(test_path)
# with open(test_path, 'r') as fid:
# test_list = [l.strip() for l in fid.readlines()]
# test_img_files = [os.path.join(dirname, 'image', f) for f in test_list]
# test_label_files = [os.path.join(dirname, 'label', f) for f in test_list]
# test_img_list = tf.data.Dataset.list_files(test_img_files,shuffle= False)
# test_label_list = tf.data.Dataset.list_files(test_label_files,shuffle= False)
# test_dataset = tf.data.Dataset.zip((test_img_list,test_label_list))
# test_dataset = test_dataset.map(load_image_test)
# test_dataset = test_dataset.batch(1)
generator = Generator()
discriminator = Discriminator()
generator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
discriminator_optimizer = tf.keras.optimizers.Adam(2e-4, beta_1=0.5)
checkpoint_best_dir = './Unet_attention_data_checkpoints/best'
checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,
discriminator_optimizer=discriminator_optimizer,
generator=generator,
discriminator=discriminator)
checkpoint.restore(tf.train.latest_checkpoint(checkpoint_best_dir))
test_psnr = 0
test_ssim = 0
for n, (input_image, target) in test_dataset.enumerate():
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
else:
dataloader = DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
dataloader = inf_train_gen(dataloader)
#models
generator = Generator().to(device)
discriminator = Discriminator().to(device)
vgg = Vgg16(requires_grad=False).to(device)
if args.pre_train:
if args.distributed:
g_checkpoint = torch.load(
args.checkpoint_path +
'generator_checkpoint_{}.ckpt'.format(args.last_iter),
map_location=lambda storage, loc: storage.cuda(args.local_rank
))
d_checkpoint = torch.load(
args.checkpoint_path +
'discriminator_checkpoint_{}.ckpt'.format(args.last_iter),
map_location=lambda storage, loc: storage.cuda(args.local_rank
))
def main():
# Load the data
data = GANstronomyDataset(opts.DATA_PATH, split=opts.TVT_SPLIT)
data.set_split_index(0)
data_loader = torch.utils.data.DataLoader(data,
batch_size=opts.BATCH_SIZE,
shuffle=True)
num_classes = data.num_classes()
# Make the output directory
util.create_dir(opts.RUN_PATH)
util.create_dir(opts.IMG_OUT_PATH)
util.create_dir(opts.MODEL_OUT_PATH)
# Copy opts.py and model.py to opts.RUN_PATH as a record
shutil.copy2('opts.py', opts.RUN_PATH)
shutil.copy2('model.py', opts.RUN_PATH)
shutil.copy2('train.py', opts.RUN_PATH)
# Instantiate the models
G = Generator(opts.EMBED_SIZE, num_classes).to(opts.DEVICE)
G_optimizer = torch.optim.Adam(G.parameters(),
lr=opts.ADAM_LR,
betas=opts.ADAM_B)
D = Discriminator(num_classes).to(opts.DEVICE)
D_optimizer = torch.optim.Adam(D.parameters(),
lr=opts.ADAM_LR,
betas=opts.ADAM_B)
if opts.MODEL_PATH is None:
start_iepoch, start_ibatch = 0, 0
else:
print('Attempting to resume training using model in %s...' %
opts.MODEL_PATH)
start_iepoch, start_ibatch = load_state_dicts(opts.MODEL_PATH, G,
G_optimizer, D,
D_optimizer)
for iepoch in range(opts.NUM_EPOCHS):
for ibatch, data_batch in enumerate(data_loader):
# To try to resume training, just continue if iepoch and ibatch are less than their starts
if iepoch < start_iepoch or (iepoch == start_iepoch
and ibatch < start_ibatch):
if iepoch % opts.INTV_PRINT_LOSS == 0 and not ibatch:
print('Skipping epoch %d...' % iepoch)
continue
recipe_ids, recipe_embs, img_ids, imgs, classes, noisy_real, noisy_fake = data_batch
# Make sure we're not training on validation or test data!
if opts.SAFETY_MODE:
for recipe_id in recipe_ids:
assert data.get_recipe_split_index(recipe_id) == 0
batch_size, recipe_embs, imgs, classes, classes_one_hot = util.get_variables(
recipe_ids, recipe_embs, img_ids, imgs, classes, num_classes)
noisy_real, noisy_fake = util.get_variables2(
noisy_real, noisy_fake)
# Adversarial ground truths
all_real = Variable(FloatTensor(batch_size, 1).fill_(1.0),
requires_grad=False).to(opts.DEVICE)
all_fake = Variable(FloatTensor(batch_size, 1).fill_(0.0),
requires_grad=False).to(opts.DEVICE)
# Train Discriminator
imgs_gen = G(recipe_embs, classes_one_hot).detach()
for _ in range(opts.NUM_UPDATE_D):
D_optimizer.zero_grad()
fake_probs = D(imgs_gen, classes_one_hot)
real_probs = D(imgs, classes_one_hot)
D_loss = (
BCELoss(fake_probs,
noisy_fake if opts.NOISY_LABELS else all_fake) +
BCELoss(real_probs,
noisy_real if opts.NOISY_LABELS else all_real))
D_loss.backward()
D_optimizer.step()
# Train Generator
for _ in range(opts.NUM_UPDATE_G):
G_optimizer.zero_grad()
imgs_gen = G(recipe_embs, classes_one_hot)
fake_probs = D(imgs_gen, classes_one_hot)
G_BCE_loss = BCELoss(fake_probs, all_real)
G_MSE_loss = MSELoss(imgs_gen, imgs)
G_loss = opts.A_BCE * G_BCE_loss + opts.A_MSE * G_MSE_loss
G_loss.backward()
G_optimizer.step()
if iepoch % opts.INTV_PRINT_LOSS == 0 and not ibatch:
print_loss(G_BCE_loss, G_MSE_loss, D_loss, iepoch)
if iepoch % opts.INTV_SAVE_IMG == 0 and not ibatch:
# Save a training image
get_img_gen(data, 0, G, iepoch, opts.IMG_OUT_PATH)
# Save a validation image
get_img_gen(data, 1, G, iepoch, opts.IMG_OUT_PATH)
if iepoch % opts.INTV_SAVE_MODEL == 0 and not ibatch:
print('Saving model...')
save_model(G, G_optimizer, D, D_optimizer, iepoch,
opts.MODEL_OUT_PATH)
save_model(G, G_optimizer, D, D_optimizer, 'FINAL', opts.MODEL_OUT_PATH)
print('\a') # Ring the bell to alert the human
def train(epochs, batchsize, iterations, data_path, modeldir, nc_size):
# Dataset definition
dataset = DatasetLoader(data_path, nc_size)
# Model & Optimizer definition
generator = Generator(nc_size)
generator.to_gpu()
gen_opt = set_optimizer(generator, 0.00005, 0.9)
discriminator = Discriminator(nc_size)
discriminator.to_gpu()
dis_opt = set_optimizer(discriminator, 0.00005, 0.9)
for epoch in range(epochs):
sum_dis_loss = 0
sum_gen_loss = 0
for batch in range(0, iterations, batchsize):
x, x_label, y, y_label, z, z_label = dataset.train(batchsize)
# Discriminator update
# Adversairal loss
a = y_label - x_label
fake = generator(x, a)
fake.unchain_backward()
loss = adversarial_loss_dis(discriminator, fake, y)
# Interpolation loss
rnd = np.random.randint(2)
if rnd == 0:
alpha = xp.random.uniform(0, 0.5, size=batchsize)
else:
alpha = xp.random.uniform(0.5, 1.0, size=batchsize)
alpha = chainer.as_variable(alpha.astype(xp.float32))
alpha = F.tile(F.expand_dims(alpha, axis=1), (1, nc_size))
fake_0 = generator(x, y_label - y_label)
fake_1 = generator(x, alpha * a)
fake_0.unchain_backward()
fake_1.unchain_backward()
loss += 10 * interpolation_loss_dis(discriminator, fake_0, fake,
fake_1, alpha, rnd)
# Matching loss
v2 = y_label - z_label
v3 = z_label - x_label
loss += matching_loss_dis(discriminator, x, fake, y, z, a, v2, v3)
discriminator.cleargrads()
loss.backward()
dis_opt.update()
loss.unchain_backward()
sum_dis_loss += loss.data
# Generator update
# Adversarial loss
fake = generator(x, a)
loss = adversarial_loss_gen(discriminator, fake)
# Interpolation loss
rnd = np.random.randint(2)
if rnd == 0:
alpha = xp.random.uniform(0, 0.5, size=batchsize)
else:
alpha = xp.random.uniform(0.5, 1.0, size=batchsize)
alpha = chainer.as_variable(alpha.astype(xp.float32))
alpha = F.tile(F.expand_dims(alpha, axis=1), (1, nc_size))
fake_alpha = generator(x, alpha * a)
loss += 10 * interpolation_loss_gen(discriminator, fake_alpha)
# Matching loss
loss += matching_loss_gen(discriminator, x, fake, a)
# Cycle-consistency loss
cyc = generator(fake, -a)
loss += 10 * F.mean_absolute_error(cyc, x)
# Self-reconstruction loss
fake_0 = generator(x, y_label - y_label)
loss += 10 * F.mean_absolute_error(fake_0, x)
generator.cleargrads()
loss.backward()
gen_opt.update()
loss.unchain_backward()
sum_gen_loss += loss.data
if batch == 0:
serializers.save_npz(f"{modeldir}/generator_{epoch}.model",
generator)
print(
f"epoch: {epoch} disloss: {sum_dis_loss/iterations} genloss: {sum_gen_loss/iterations}"
)
dataset = Cifar100()
stm = PreResNet(depth=32, num_classes=100)
ltm = PreResNet(depth=32, num_classes=100)
seg = []
for module in ltm.modules():
if isinstance(module, nn.Conv2d) and module.weight.data.shape[
1] % args.units_x == 0 and module.weight.data.shape[2] == 3:
seg.append(module.weight.data.shape[0] // args.units_x)
print("seg: ", seg)
dim_output = 64 # output dim of PreResNet
gen = Generator(in_features=dim_output + args.hidden_dim,
out_features=sum(seg) * args.num_units,
total_class=args.total_class,
hidden_dim=args.hidden_dim,
num_units=args.num_units,
units_x=args.units_x,
units_y=args.units_y,
seg=seg)
units = torch.rand(args.num_units, args.units_x, args.units_y).unsqueeze(0)
units = units.repeat(args.batch_size, 1, 1, 1)
# create trainer
trainer = ModelTrainer(stm=stm,
ltm=ltm,
gen=gen,
dataset=dataset,
units=units)
trainer.train(batch_size=args.batch_size, epoches=args.epoches, lr=args.lr)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--input_dir', help = 'Directory containing xxx_i_s and xxx_i_t with same prefix',
default = cfg.example_data_dir)
parser.add_argument('--save_dir', help = 'Directory to save result', default = cfg.predict_result_dir)
parser.add_argument('--checkpoint', help = 'ckpt', default = cfg.ckpt_path)
args = parser.parse_args()
assert args.input_dir is not None
assert args.save_dir is not None
assert args.checkpoint is not None
print_log('model compiling start.', content_color = PrintColor['yellow'])
G = Generator(in_channels = 3).to(device)
D1 = Discriminator(in_channels = 6).to(device)
D2 = Discriminator(in_channels = 6).to(device)
vgg_features = Vgg19().to(device)
G_solver = torch.optim.Adam(G.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2))
D1_solver = torch.optim.Adam(D1.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2))
D2_solver = torch.optim.Adam(D2.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2))
checkpoint = torch.load(args.checkpoint)
G.load_state_dict(checkpoint['generator'])
D1.load_state_dict(checkpoint['discriminator1'])
D2.load_state_dict(checkpoint['discriminator2'])
G_solver.load_state_dict(checkpoint['g_optimizer'])
D1_solver.load_state_dict(checkpoint['d1_optimizer'])
D2_solver.load_state_dict(checkpoint['d2_optimizer'])
trfms = To_tensor()
example_data = example_dataset(data_dir= args.input_dir, transform = trfms)
example_loader = DataLoader(dataset = example_data, batch_size = 1, shuffle = False)
example_iter = iter(example_loader)
print_log('Model compiled.', content_color = PrintColor['yellow'])
print_log('Predicting', content_color = PrintColor['yellow'])
with torch.no_grad():
for step in tqdm(range(len(example_data)//2)):
try:
inp = example_iter.next()
except StopIteration:
example_iter = iter(example_loader)
inp = example_iter.next()
i_t = inp[0].to(device)
i_s = inp[1].to(device)
name = str(inp[2][0])
o_sk, o_t, o_b, o_f = G(i_t, i_s)
o_sk = o_sk.squeeze(0).detach().to('cpu')
o_t = o_t.squeeze(0).detach().to('cpu')
o_b = o_b.squeeze(0).detach().to('cpu')
o_f = o_f.squeeze(0).detach().to('cpu')
if not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
o_sk = F.to_pil_image((o_sk)*255.0)
o_t = F.to_pil_image((o_t + 1)*127.5)
o_b = F.to_pil_image((o_b + 1)*127.5)
o_f = F.to_pil_image((o_f + 1)*127.5)
o_f.save(os.path.join(args.save_dir, name + 'o_f.png'))
num_gpus = len(gpus)
parser = argparse.ArgumentParser(description='Progressive Growing of GANs')
parser.add_argument('path',
type=str,
help='path of specified dataset',
default='/raid/datasets/img_align_celeba')
parser.add_argument('-d',
'--data',
default='celeba',
type=str,
choices=['celeba', 'lsun'],
help=('Specify dataset. '
'Currently CelebA and LSUN is supported'))
generator = Generator(code_size, n_label).cuda()
discriminator = Discriminator(n_label).cuda()
g_running = Generator(code_size, n_label).cuda()
g_running.train(False)
generator = torch.nn.DataParallel(generator, device_ids=gpus)
discriminator = torch.nn.DataParallel(discriminator, device_ids=gpus)
g_running = torch.nn.DataParallel(g_running, device_ids=gpus)
class_loss = nn.CrossEntropyLoss()
g_optimizer = optim.Adam(generator.parameters(), lr=0.001, betas=(0.0, 0.99))
d_optimizer = optim.Adam(discriminator.parameters(),
lr=0.001,
betas=(0.0, 0.99))
type=str,
default="factor",
help="filename prefix to result samples",
)
parser.add_argument(
"factor",
type=str,
help="name of the closed form factorization result factor file",
)
args = parser.parse_args()
eigvec = torch.load(args.factor)["eigvec"].to(args.device)
ckpt = torch.load(args.ckpt)
g = Generator(args.size,
512,
8,
channel_multiplier=args.channel_multiplier).to(args.device)
g.load_state_dict(ckpt["g_ema"], strict=False)
trunc = g.mean_latent(4096)
for i in range(1, 512):
for j in range(1, 50):
latent = torch.randn(args.n_sample, 512, device=args.device)
latent = g.get_latent(latent)
direction = j * eigvec[:, i].unsqueeze(0)
img, _ = g(
[latent],
truncation=args.truncation,
truncation_latent=trunc,
MODEL_NAME = "./pretrained/checkpoints/checkpoint_13399.ckpt"
GENERATE_N_SAMPLES = 100000
OUTPUT_DIR = "./samples/generate.txt"
BATCH_SIZE = 64
SEQ_LEN = 10
DIM = 128
with open("pretrained/char2int.pickle", "rb") as f:
charmap = pickle.load(f)
f.close()
with open("pretrained/chars.pickle", "rb") as f:
chars = pickle.load(f)
f.close()
fake_inputs = Generator(BATCH_SIZE, SEQ_LEN, DIM, len(charmap))
with tf.Session() as session:
def generate_samples():
samples = session.run(fake_inputs)
samples = np.argmax(samples, axis=2)
decoded_samples = []
for i in range(len(samples)):
decoded = []
for j in range(len(samples[i])):
decoded.append(chars[samples[i][j]])
decoded_samples.append(tuple(decoded))
return decoded_samples
def save(samples):
parser.add_argument('--upscale_factor', default=1, type=int, choices=[1, 4, 8],
help='super resolution upscale factor')
parser.add_argument('--num_epochs', default=1, type=int, help='train epoch number')
opt = parser.parse_args()
CROP_SIZE = opt.crop_size #裁剪会带来拼尽问题嘛
UPSCALE_FACTOR = opt.upscale_factor #上采样
NUM_EPOCHS = opt.num_epochs #轮数
val_set = TestDatasetFromFolder('/data/lpw/FusionDataset/tmp_val/', upscale_factor=UPSCALE_FACTOR) #测试集导入
for pthdir in os.listdir('/data/lpw/ResnetFusion/epochs/'):
for MODEL_NAME in os.listdir('/data/lpw/ResnetFusion/epochs/'+pthdir):
# MODEL_NAME = 'netG_epoch_1_4000.pth'
netG = Generator(UPSCALE_FACTOR).eval()
netG.cuda()
netG.load_state_dict(torch.load('/data/lpw/ResnetFusion/epochs/' + pthdir+'/'+MODEL_NAME))
val_loader = DataLoader(dataset=val_set, num_workers=1, batch_size=1, shuffle=False)
epoch =1
out_path = '/data/lpw/ResnetFusion/quota_results/SRF_' + str(UPSCALE_FACTOR) + '/'+pthdir+'/'+MODEL_NAME+'/'#输出路径
print(out_path)
if not os.path.exists(out_path):
os.makedirs(out_path)
val_bar = tqdm(val_loader) #验证集的进度条
val_images = []
index = 1
for val_lr , val_lr_restore, val_hr in val_bar:
batch_size = val_lr.size(0)
def train(args):
device = torch.device('cuda' if torch.cuda.is_available() and args.enable_cuda else 'cpu')
# transforms applied
transform = transforms.Compose([
transforms.Resize((args.image_size, args.image_size)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
if args.dataset == 'celeba':
img_dir = 'cropped'
ann_dir = 'list_attr_celeba.csv'
train_dataset = CelebA(args.root_dir, img_dir, ann_dir, transform=transform, train=True)
test_dataset = CelebA(args.root_dir, img_dir, ann_dir, transform=transform, train=False)
attnames = list(train_dataset.df.columns)[1:]
elif args.dataset == 'sunattributes':
img_dir = 'cropped'
ann_dir = 'SUNAttributeDB'
train_dataset = SUN_Attributes(args.root_dir, img_dir, ann_dir, transform=transform, train=True)
test_dataset = SUN_Attributes(args.root_dir, img_dir, ann_dir, transform=transform, train=False)
attnames = train_dataset.attrnames
else:
raise NotImplementedError()
fsize = train_dataset.feature_size
trainloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn, drop_last=True)
testloader = DataLoader(test_dataset, batch_size=args.show_size, shuffle=True, collate_fn=collate_fn, drop_last=True)
'''
dataloader returns dictionaries.
sample : {'image' : (bs, 64, 64, 3), 'attributes' : (bs, att_size)}
'''
# model, optimizer, criterion
if args.residual:
gen = Generator_Res(in_c = args.nz + fsize)
dis = Discriminator_Res(ny=fsize)
else:
gen = Generator(in_c = args.nz + fsize)
dis = Discriminator(ny=fsize)
MODELPATH = '../model/celeba/res_False/gen_epoch_{}.ckpt'.format(args.model_ep)
gen.load_state_dict(torch.load(MODELPATH))
enc_y = Encoder(fsize).to(device)
MODELPATH = '../model/celeba/res_False/enc_y_epoch_{}.ckpt'.format(args.enc_ep)
enc_y.load_state_dict(torch.load(MODELPATH))
enc_z = Encoder(args.nz, for_y=False).to(device)
MODELPATH = '../model/celeba/res_False/enc_y_epoch_{}.ckpt'.format(args.enc_ep)
enc_z.load_state_dict(torch.load(MODELPATH))
gen.eval()
enc_y.eval()
enc_z.eval()
model = AttrEncoder(outdims=fsize).to(device)
# initialize weights for encoders
att_optim = optim.Adam(model.parameters(), lr=args.learning_rate, betas=args.betas)
criterion = nn.BCELoss().to(device)
noise = torch.randn((args.batch_size, args.nz)).to(device)
if args.use_tensorboard:
writer.add_text("Text", "begin training, lr={}".format(args.learning_rate))
print("begin training, lr={}".format(args.learning_rate), flush=True)
stepcnt = 0
for ep in range(args.num_epoch):
for it, sample in enumerate(trainloader):
elapsed = time.time()
image = sample['image'].to(device)
att = sample['attributes'].to(device)
out = model(image)
loss = criterion(out, att)
z = enc_z(image)
y = enc_y(image)
out2 = gen(z, y)
loss2 = criterion(out2, att)
loss.backward()
loss2.backward()
att_optim.step()
if it % args.log_every == (args.log_every - 1):
if args.use_tensorboard:
writer.add_scalar('loss', loss, it+1)
print("{}th iter \t loss: {:.8f} \t time per iter: {:.05f}s".format(it+1, loss.detach().cpu(), (time.time() - elapsed) / args.log_every), flush=True)
cnt, allcnt = eval(model, testloader, device)
print("-" * 50)
print("epoch {} done. accuracy: {:.03f}%. num guessed: [{:05d}/{:05d}]".format(ep+1, cnt / allcnt * 100, cnt, allcnt))
print("-" * 50, flush=True)
if ep % args.save_every == (args.save_every - 1):
torch.save(model.state_dict(), "../model/atteval_epoch_{}.model".format(ep+1))
cnt, allcnt = eval(model, testloader, device, early=False)
print("-" * 50)
print("training done. final accuracy: {:.03f}% num guessed: [{:07d}/{:07d}]".format(ep+1, cnt / allcnt * 100, cnt, allcnt))
print("-" * 50, flush=True)
def main(_):
strategy = tf.distribute.MirroredStrategy()
NUM_GPU = len(tf.config.experimental.list_physical_devices('GPU'))
train_ds, ds_info = tfds.load(FLAGS.dataset,
split='train',
shuffle_files=True,
with_info=True)
#dataset is very big, don't want to wait long
if FLAGS.dataset == 'lsun/bedroom':
train_ds = train_ds.take(300000)
output_channels = 3
if FLAGS.dataset == 'cifar10':
output_channels = 3
if FLAGS.dataset == 'mnist':
output_channels = 1
OUTPUT_DIM = FLAGS.image_size * FLAGS.image_size * output_channels
def preprocess(image):
"""Normalize the images to [-1.0, 1.0]"""
image = image['image']
image = tf.image.resize_with_pad(image, FLAGS.image_size,
FLAGS.image_size)
return (tf.cast(image, tf.float32) - 127.5) / 127.5
train_ds = train_ds.map(preprocess,
num_parallel_calls=tf.data.experimental.AUTOTUNE)
train_ds = train_ds.cache()
train_ds = train_ds.shuffle(ds_info.splits['train'].num_examples)
train_ds = train_ds.batch(FLAGS.batch_size)
train_ds = train_ds.prefetch(tf.data.experimental.AUTOTUNE)
train_ds = strategy.experimental_distribute_dataset(train_ds)
with strategy.scope():
generator_optimizer = tf.keras.optimizers.Adam(learning_rate=FLAGS.lr,
beta_1=0.5,
beta_2=0.9)
discriminator_optimizer = tf.keras.optimizers.Adam(
learning_rate=FLAGS.lr, beta_1=0.5, beta_2=0.9)
inputs = tf.keras.Input(shape=(FLAGS.latent_vector, ),
name="latent_vector")
outputs = Generator(FLAGS.num_filters, FLAGS.latent_vector,
output_channels, OUTPUT_DIM)(inputs)
generator = tf.keras.Model(inputs=inputs, outputs=outputs)
inputs = tf.keras.Input(shape=(FLAGS.image_size * FLAGS.image_size *
output_channels),
name="imgs")
outputs = Discriminator(FLAGS.num_filters, FLAGS.image_size,
output_channels)(inputs)
discriminator = tf.keras.Model(inputs=inputs, outputs=outputs)
@tf.function
def train_gen():
noise = tf.random.normal(
[FLAGS.batch_size // NUM_GPU, FLAGS.latent_vector])
with tf.GradientTape() as gen_tape:
generated_images = generator(noise, training=True)
fake_output = discriminator(tf.reshape(generated_images,
[-1, OUTPUT_DIM]),
training=False)
gen_loss = -tf.reduce_mean(fake_output)
gradients_of_generator = gen_tape.gradient(
gen_loss, generator.trainable_variables)
generator_optimizer.apply_gradients(
zip(gradients_of_generator, generator.trainable_variables))
return tf.reduce_mean(gen_loss)
@tf.function
def train_disc(images):
image = tf.reshape(images, [-1, OUTPUT_DIM])
noise = tf.random.normal([images.shape[0], FLAGS.latent_vector])
fake_images = generator(noise, training=True)
with tf.GradientTape() as disc_tape:
disc_real = discriminator(images, training=True)
disc_fake = discriminator(fake_images, training=True)
disc_loss = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)
alpha = tf.random.uniform(shape=[
images.shape[0],
1,
],
minval=0.,
maxval=1.)
differences = fake_images - image
interpolates = image + (alpha * differences)
gradients = tf.gradients(discriminator(interpolates),
[interpolates])[0]
slopes = tf.math.sqrt(tf.reduce_sum(tf.square(gradients),
axis=[1]))
gradient_penalty = tf.reduce_mean((slopes - 1.)**2)
disc_loss += 10 * gradient_penalty
gradients_of_discriminator = disc_tape.gradient(
disc_loss, discriminator.trainable_variables)
discriminator_optimizer.apply_gradients(
zip(gradients_of_discriminator, discriminator.trainable_variables))
return tf.reduce_mean(disc_loss)
@tf.function
def distributed_disc_step(dist_inputs):
per_replica_disc_loss = strategy.run(train_disc, args=[dist_inputs])
return strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_disc_loss,
axis=None)
@tf.function
def distributed_gen_step():
per_replica_gen_loss = strategy.run(train_gen, args=())
return strategy.reduce(tf.distribute.ReduceOp.SUM,
per_replica_gen_loss,
axis=None)
def save_images(model, ep, vector):
predictions = tf.clip_by_value(model(vector, training=False), -1, 1)
plt.figure(figsize=(5, 5))
for i in range(predictions.shape[0]):
plt.subplot(4, 4, i + 1)
pred = tf.reshape(
predictions[i],
[FLAGS.image_size, FLAGS.image_size, output_channels])
plt.imshow((pred.numpy() * 127.5 + 127.5).astype(np.uint8))
plt.axis('off')
plt.savefig(FLAGS.save_folder +
'/image_at_epoch_{:02d}.png'.format(ep))
if not os.path.exists(FLAGS.save_folder):
os.makedirs(FLAGS.save_folder)
noise_vector = tf.random.normal([FLAGS.num_examples, FLAGS.latent_vector])
for epoch in tqdm(range(FLAGS.epochs)):
iterator = iter(train_ds)
gen_loss = 0
disc_loss = 0
num_batch = 0
iterations = 0
flag = True
while flag:
gen_loss += distributed_gen_step()
iterations += 1
for _ in range(FLAGS.disc_iters):
optional = iterator.get_next_as_optional()
if optional.has_value().numpy() == False:
flag = False
else:
data = optional.get_value()
d_loss = distributed_disc_step(data)
disc_loss += d_loss
num_batch += 1
disc_loss /= num_batch
gen_loss /= iterations
print("Epoch {}, gen_loss {:.5f} \n disc_loss {:.5f}\n".format(
epoch, gen_loss, disc_loss))
save_images(generator, epoch, noise_vector)
save_images(generator, FLAGS.epochs, noise_vector)
def convert(config):
os.makedirs(join(config.convert_dir, config.resume_model), exist_ok=True)
sampling_rate, num_mcep, frame_period = config.sampling_rate, 36, 5
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Restore model
print(f'Loading the trained models from step {config.resume_model}...')
generator = Generator().to(device)
g_path = join(config.model_save_dir, f'{config.resume_model}-G.ckpt')
generator.load_state_dict(
torch.load(g_path, map_location=lambda storage, loc: storage))
# for all possible speaker pairs in config.speakers
for i in range(0, len(config.speakers)):
for j in range(0, len(config.speakers)):
if i != j:
target_dir = join(
config.convert_dir, str(config.resume_model),
f'{config.speakers[i]}_to_{config.speakers[j]}')
os.makedirs(target_dir, exist_ok=True)
# Load speakers
data_loader = ConvertDataset(config,
src_spk=config.speakers[i],
trg_spk=config.speakers[j])
print('---------------------------------------')
print('Source: ', config.speakers[i], ' Target: ',
config.speakers[j])
print('---------------------------------------')
# Read a batch of testdata
src_test_wavfiles = data_loader.get_batch_test_data(
batch_size=config.num_converted_wavs)
src_test_wavs = [
load_wav(wavfile, sampling_rate)
for wavfile in src_test_wavfiles
]
with torch.no_grad():
for idx, wav in enumerate(src_test_wavs):
print(f'({idx}), file length: {len(wav)}')
wav_name = basename(src_test_wavfiles[idx])
# convert wav to mceps
f0, _, sp, ap = world_decompose(
wav=wav,
fs=sampling_rate,
frame_period=frame_period)
f0_converted = pitch_conversion(
f0=f0,
mean_log_src=data_loader.logf0s_mean_src,
std_log_src=data_loader.logf0s_std_src,
mean_log_target=data_loader.logf0s_mean_trg,
std_log_target=data_loader.logf0s_std_trg)
coded_sp = world_encode_spectral_envelop(
sp=sp, fs=sampling_rate, dim=num_mcep)
print("Before being fed into G: ", coded_sp.shape)
coded_sp_norm = (coded_sp - data_loader.mcep_mean_src
) / data_loader.mcep_std_src
coded_sp_norm_tensor = torch.FloatTensor(
coded_sp_norm.T).unsqueeze_(0).unsqueeze_(1).to(
device)
spk_conds = torch.FloatTensor(
data_loader.spk_c_trg).to(device)
# generate converted speech
coded_sp_converted_norm = generator(
coded_sp_norm_tensor,
spk_conds).data.cpu().numpy()
coded_sp_converted = np.squeeze(
coded_sp_converted_norm
).T * data_loader.mcep_std_trg + data_loader.mcep_mean_trg
coded_sp_converted = np.ascontiguousarray(
coded_sp_converted)
print("After being fed into G: ",
coded_sp_converted.shape)
# convert back to wav
wav_transformed = world_speech_synthesis(
f0=f0_converted,
coded_sp=coded_sp_converted,
ap=ap,
fs=sampling_rate,
frame_period=frame_period)
wav_id = wav_name.split('.')[0]
# SAVE TARGET SYNTHESIZED
librosa.output.write_wav(
join(target_dir,
f'{wav_id}-vcto-{data_loader.trg_spk}.wav'),
wav_transformed, sampling_rate)
# SAVE COPY OF TARGET REFERENCE
wav_num = wav_name.split('.')[0].split('_')[1]
copy(
f'{config.wav_dir}/{config.speakers[j]}/{config.speakers[j]}_{wav_num}.wav',
target_dir)
args = parser.parse_args()
sys.path.append(args.repo)
import dnnlib
from dnnlib import tflib
tflib.init_tf()
with open(args.path, 'rb') as f:
generator, discriminator, g_ema = pickle.load(f)
size = g_ema.output_shape[2]
g = Generator(size, 512, 8, channel_multiplier=args.channel_multiplier)
state_dict = g.state_dict()
state_dict = fill_statedict(state_dict, g_ema.vars, size)
g.load_state_dict(state_dict)
latent_avg = torch.from_numpy(g_ema.vars['dlatent_avg'].value().eval())
ckpt = {'g_ema': state_dict, 'latent_avg': latent_avg}
if args.gen:
g_train = Generator(size,
512,
8,
channel_multiplier=args.channel_multiplier)
g_train_state = g_train.state_dict()
def main():
parser = argparse.ArgumentParser()
parser.add_argument('--batchsize', '-b', type=int, default=12)
parser.add_argument('--epoch', '-e', type=int, default=500)
parser.add_argument('--gpu', '-g', type=int, default=-1)
parser.add_argument('--root', '-R', default='/mnt/sakuradata10-striped/gao/background')
parser.add_argument('--out', '-o', default='/mnt/sakuradata10-striped/gao/results/cartoongan')
parser.add_argument('--resume', '-r', default='', help='snapshot No.')
parser.add_argument('--model_num', '-m', default='', help='generater No.')
parser.add_argument('--snapshot_interval', type=int, default=1000)
parser.add_argument('--test_interval', type=int, default=100)
parser.add_argument('--display_interval', type=int, default=5)
parser.add_argument('--size', type=int, default=256)
parser.add_argument('--use_gan', '-G', action='store_true')
args = parser.parse_args()
print('GPU: {}'.format(args.gpu))
print('# Minibatch-size: {}\n'.format(args.batchsize))
# Setup models
gen = Generator()
vgg = VGG()
if args.use_gan:
dis = Discriminator()
else:
dis = None
# Setup datasets
photos = PhotoDataset(os.path.join(args.root, 'photos_resized', '*'), crop_size=args.size)
photos_iter = chainer.iterators.SerialIterator(photos, args.batchsize)
if args.use_gan:
illusts = ImageDataset(os.path.join(args.root, 'illusts', '*', '*', '*'), crop_size=args.size)
illusts_iter = chainer.iterators.SerialIterator(illusts, args.batchsize)
else:
illusts_iter = None
# models to gpu
if args.gpu >= 0:
chainer.cuda.get_device_from_id(args.gpu).use()
gen.to_gpu()
vgg.to_gpu()
if args.use_gan:
dis.to_gpu()
# Setup optimizer parameters.
opt = chainer.optimizers.Adam(alpha=0.0002)
opt.setup(gen)
if args.use_gan:
opt_d = chainer.optimizers.Adam(alpha=0.0002)
opt_d.setup(dis)
else:
opt_d = None
# Set up a trainer
optimizers = {'gen': opt, 'dis': opt_d} if args.use_gan else {'gen': opt}
iterators = {'main': photos_iter, 'illusts': illusts_iter} if args.use_gan else {'main': photos_iter}
updater = CartoonGAN(
models=(gen, dis, vgg),
iterator=iterators,
optimizer=optimizers,
device=args.gpu,
w=10
)
out = os.path.join(args.out, 'gan') if args.use_gan else os.path.join(args.out, 'initial')
trainer = training.Trainer(updater, (args.epoch, 'epoch'), out=out)
# Load npz if necessary
if args.resume:
chainer.serializers.load_npz(os.path.join(out, 'snapshot_iter_'+args.resume+'.npz'), trainer)
print('snapshot {} loaded\n'.format(args.resume))
elif args.model_num:
chainer.serializers.load_npz(os.path.join(args.out, 'initial', 'gen_iter_'+args.model_num+'.npz'), gen)
print('model {} loaded\n'.format(args.model_num))
# trainer extensions
snapshot_interval = (args.snapshot_interval, 'iteration')
test_interval = (args.test_interval, 'iteration')
trainer.extend(extensions.dump_graph('gen/loss'))
trainer.extend(
extensions.snapshot(filename='snapshot_iter_{.updater.iteration}.npz'),
trigger=snapshot_interval)
trainer.extend(extensions.snapshot_object(
gen, 'gen_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.LogReport(trigger=(args.display_interval, 'iteration'), ))
report = ['epoch', 'iteration', 'gen/loss', 'gen/content', 'gen/mae']
if args.use_gan:
report += ['gen/adv', 'dis/illust', 'dis/edge', 'dis/photo', 'dis/loss']
trainer.extend(extensions.snapshot_object(
dis, 'dis_iter_{.updater.iteration}.npz'), trigger=snapshot_interval)
trainer.extend(extensions.PrintReport(report))
trainer.extend(extensions.ProgressBar(update_interval=args.display_interval))
trainer.extend(photos.visualizer(), trigger=test_interval)
trainer.run()
# Save the trained model
chainer.serializers.save_npz(os.path.join(args.out, 'model_final'), gen)
chainer.serializers.save_npz(os.path.join(args.out, 'optimizer_final'), opt)
if (i + 1) % 10000 == 0:
no = str(i + 1).zfill(7)
torch.save(generator.state_dict(), f'checkpoint/generator_{no}.pt')
torch.save(discriminator.state_dict(), f'checkpoint/discriminator_{no}.pt')
torch.save(g_optimizer.state_dict(), f'checkpoint/gen_optimizer_{no}.pt')
torch.save(d_optimizer.state_dict(), f'checkpoint/dis_optimizer_{no}.pt')
pbar.set_description(
(f'{i + 1}; G: {gen_loss_val:.5f};' f' D: {disc_loss_val:.5f}')
)
if __name__ == '__main__':
args = parser.parse_args()
print(args)
n_class = len(glob.glob(os.path.join(args.path, '*/')))
if args.model == 'dcgan':
from model import Generator, Discriminator
elif args.model == 'resnet':
from model_resnet import Generator, Discriminator
generator = Generator(args.code, n_class).to(device)
discriminator = Discriminator(n_class).to(device)
g_optimizer = optim.Adam(generator.parameters(), lr=args.lr_g, betas=(0, 0.9))
d_optimizer = optim.Adam(discriminator.parameters(), lr=args.lr_d, betas=(0, 0.9))
train(args, n_class, generator, discriminator)
def main():
os.environ['CUDA_VISIBLE_DEVICES'] = str(cfg.gpu)
train_name = get_train_name()
print_log('Initializing SRNET', content_color = PrintColor['yellow'])
train_data = datagen_srnet(cfg)
train_data = DataLoader(dataset = train_data, batch_size = cfg.batch_size, shuffle = False, collate_fn = custom_collate, pin_memory = True)
trfms = To_tensor()
example_data = example_dataset(transform = trfms)
example_loader = DataLoader(dataset = example_data, batch_size = len(example_data), shuffle = False)
print_log('training start.', content_color = PrintColor['yellow'])
G = Generator(in_channels = 3).cuda()
D1 = Discriminator(in_channels = 6).cuda()
D2 = Discriminator(in_channels = 6).cuda()
vgg_features = Vgg19().cuda()
G_solver = torch.optim.Adam(G.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2))
D1_solver = torch.optim.Adam(D1.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2))
D2_solver = torch.optim.Adam(D2.parameters(), lr=cfg.learning_rate, betas = (cfg.beta1, cfg.beta2))
g_scheduler = torch.optim.lr_scheduler.MultiStepLR(G_solver, milestones=[30, 200], gamma=0.5)
d1_scheduler = torch.optim.lr_scheduler.MultiStepLR(D1_solver, milestones=[30, 200], gamma=0.5)
d2_scheduler = torch.optim.lr_scheduler.MultiStepLR(D2_solver, milestones=[30, 200], gamma=0.5)
requires_grad(G, False)
requires_grad(D1, True)
requires_grad(D2, True)
disc_loss_val = 0
gen_loss_val = 0
grad_loss_val = 0
trainiter = iter(train_data)
example_iter = iter(example_loader)
K = torch.nn.ZeroPad2d((0, 1, 1, 0))
for step in tqdm(range(cfg.max_iter)):
D1_solver.zero_grad()
D2_solver.zero_grad()
if ((step+1) % cfg.save_ckpt_interval == 0):
torch.save(
{
'generator': G.module.state_dict(),
'discriminator1': D1.module.state_dict(),
'discriminator2': D2.module.state_dict(),
'g_optimizer': G_solver.state_dict(),
'd1_optimizer': D1_solver.state_dict(),
'd2_optimizer': D2_solver.state_dict(),
},
f'checkpoint/train_step-{step+1}.model',
)
i_t, i_s, t_sk, t_t, t_b, t_f, mask_t = trainiter.next()
i_t = i_t.cuda()
i_s = i_s.cuda()
t_sk = t_sk.cuda()
t_t = t_t.cuda()
t_b = t_b.cuda()
t_f = t_f.cuda()
mask_t = mask_t.cuda()
#inputs = [i_t, i_s]
labels = [t_sk, t_t, t_b, t_f]
o_sk, o_t, o_b, o_f = G(i_t, i_s)
o_sk = K(o_sk)
o_t = K(o_t)
o_b = K(o_b)
o_f = K(o_f)
#print(o_sk.shape, o_t.shape, o_b.shape, o_f.shape)
#print('------')
#print(i_s.shape)
i_db_true = torch.cat((t_b, i_s), dim = 1)
i_db_pred = torch.cat((o_b, i_s), dim = 1)
i_df_true = torch.cat((t_f, i_t), dim = 1)
i_df_pred = torch.cat((o_f, i_t), dim = 1)
o_db_true = D1(i_db_true)
o_db_pred = D1(i_db_pred)
o_df_true = D2(i_df_true)
o_df_pred = D2(i_df_pred)
i_vgg = torch.cat((t_f, o_f), dim = 0)
out_vgg = vgg_features(i_vgg)
db_loss = build_discriminator_loss(o_db_true, o_db_pred)
df_loss = build_discriminator_loss(o_df_true, o_df_pred)
db_loss.backward()
df_loss.backward()
D1_solver.step()
D2_solver.step()
d1_scheduler.step()
d2_scheduler.step()
clip_grad(D1)
clip_grad(D2)
if ((step+1) % 5 == 0):
requires_grad(G, True)
requires_grad(D1, False)
requires_grad(D2, False)
G_solver.zero_grad()
o_sk, o_t, o_b, o_f = G(i_t, i_s)
o_sk = K(o_sk)
o_t = K(o_t)
o_b = K(o_b)
o_f = K(o_f)
#print(o_sk.shape, o_t.shape, o_b.shape, o_f.shape)
#print('------')
#print(i_s.shape)
i_db_true = torch.cat((t_b, i_s), dim = 1)
i_db_pred = torch.cat((o_b, i_s), dim = 1)
i_df_true = torch.cat((t_f, i_t), dim = 1)
i_df_pred = torch.cat((o_f, i_t), dim = 1)
o_db_pred = D1(i_db_pred)
o_df_pred = D2(i_df_pred)
i_vgg = torch.cat((t_f, o_f), dim = 0)
out_vgg = vgg_features(i_vgg)
out_g = [o_sk, o_t, o_b, o_f, mask_t]
out_d = [o_db_pred, o_df_pred]
g_loss, detail = build_generator_loss(out_g, out_d, out_vgg, labels)
g_loss.backward()
G_solver.step()
g_scheduler.step()
requires_grad(G, False)
requires_grad(D1, True)
requires_grad(D2, True)
if ((step+1) % cfg.write_log_interval == 0):
print('Iter: {}/{} | Gen: {} | D_bg: {} | D_fus: {}'.format(step+1, cfg.max_iter, g_loss.item(), db_loss.item(), df_loss.item()))
if ((step+1) % cfg.gen_example_interval == 0):
savedir = os.path.join(cfg.example_result_dir, train_name, 'iter-' + str(step+1).zfill(len(str(cfg.max_iter))))
with torch.no_grad():
inp = example_iter.next()
o_sk, o_t, o_b, o_f = G(inp)
if not os.path.exists(save_dir):
os.makedirs(save_dir)
o_sk = skimage.img_as_ubyte(o_sk)
o_t = skimage.img_as_ubyte(o_t + 1)
o_b = skimage.img_as_ubyte(o_b + 1)
o_f = skimage.img_as_ubyte(o_f + 1)
io.imsave(os.path.join(save_dir, name + 'o_f.png'), o_f)
io.imsave(os.path.join(save_dir, name + 'o_sk.png'), o_sk)
io.imsave(os.path.join(save_dir, name + 'o_t.png'), o_t)
io.imsave(os.path.join(save_dir, name + 'o_b.png'), o_b)
if cuda:
label_real = label_real.cuda()
label_fake = label_fake.cuda()
label_gen_real = label_gen_real.cuda()
label_gen_fake = label_gen_fake.cuda()
dis_loss = criterion(real.view(-1, 1), label_real) + criterion(
fake.view(-1, 1), label_fake)
gen_loss = criterion(fake.view(-1, 1), label_gen_real) + criterion(
real.view(-1, 1), label_gen_fake)
return dis_loss, gen_loss
if __name__ == '__main__':
A_tr, A_te, B_tr, B_te = data_read('./data', 'horse', 'zebra', IMAGE_SIZE,
BATCH_SIZE, True)
generatorA = Generator()
generatorB = Generator()
discriminatorA = Discriminator()
discriminatorB = Discriminator()
if CUDA:
generatorA = generatorA.cuda()
generatorB = generatorB.cuda()
discriminatorA = discriminatorA.cuda()
discriminatorB = discriminatorB.cuda()
gen_params = chain(generatorA.parameters(), generatorB.parameters())
dis_params = chain(discriminatorA.parameters(),
discriminatorB.parameters())
optim_gen = optim.Adam(gen_params,
lr=LR,
betas=(0.5, 0.999),
weight_decay=1e-5)
download=True,
transform=transform)
testset = torchvision.datasets.MNIST(root=download_path,
train=False,
download=True,
transform=transform)
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=BATCH * num_gpus,
shuffle=True,
num_workers=64)
testloader = torch.utils.data.DataLoader(testset,
batch_size=BATCH * num_gpus,
shuffle=False,
num_workers=64)
G = Generator()
C = Critic()
if torch.cuda.is_available():
G = G.cuda()
C = C.cuda()
if torch.cuda.device_count() > 1:
G = torch.nn.DataParallel(G)
C = torch.nn.DataParallel(C)
if cp > 0:
load_G = 'model_parameters/G' + str(cp) + '.pth'
load_C = 'model_parameters/C' + str(cp) + '.pth'
state_dict_G = torch.load(load_G)
state_dict_C = torch.load(load_C)
G.load_state_dict(state_dict_G["G_state_dict:"])
args.distributed = n_gpu > 1
if args.distributed:
torch.cuda.set_device(args.local_rank)
torch.distributed.init_process_group(backend="nccl",
init_method="env://")
synchronize()
args.latent = 512
args.n_mlp = 8
args.start_iter = 0
generator = Generator(
args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
discriminator = Discriminator(
args.size, channel_multiplier=args.channel_multiplier).to(device)
g_ema = Generator(args.size,
args.latent,
args.n_mlp,
channel_multiplier=args.channel_multiplier).to(device)
g_ema.eval()
accumulate(g_ema, generator, 0)
g_reg_ratio = args.g_reg_every / (args.g_reg_every + 1)
d_reg_ratio = args.d_reg_every / (args.d_reg_every + 1)
g_optim = optim.Adam(
utils.plot_lines(points=dset.p,
title='Weight of each gaussian',
path='{}/gaussian_weight.png'.format(exp_dir))
sample_points = dset.sample(100)
utils.plot_scatter(points=sample_points,
centers=dset.centers,
title='Sampled data points',
path='{}/samples.png'.format(exp_dir))
prefix = "unrolled_steps-{}-prior_std-{:.2f}".format(
config.unrolled_steps, np.std(dset.p))
print("Save file with prefix", prefix)
G = Generator(input_size=config.g_inp,
hidden_size=config.g_hid,
output_size=config.g_out).cuda()
G._apply(lambda t: t.detach().checkpoint())
D = Discriminator(input_size=config.d_inp,
hidden_size=config.d_hid,
output_size=config.d_out).cuda()
D._apply(lambda t: t.detach().checkpoint())
criterion = nn.BCELoss(
) # Binary cross entropy: http://pytorch.org/docs/nn.html#bceloss
def binary_cross_entropy(x, y):
loss = -(x.log() * y + (1 - x).log() * (1 - y))
return loss.mean()
criterion = binary_cross_entropy
help='using GPU or CPU')
parser.add_argument('--image_name',
type=str,
help='test low resolution image name')
parser.add_argument('--weights',
default='weights/netG_epoch_4_54.pth',
type=str,
help='generator model epoch name')
opt = parser.parse_args()
UPSCALE_FACTOR = opt.upscale_factor
TEST_MODE = True if opt.test_mode == 'GPU' else False
IMAGE_NAME = opt.image_name
weights_path = opt.weights
model = Generator(UPSCALE_FACTOR).eval()
if TEST_MODE:
model.cuda()
model.load_state_dict(torch.load(weights_path))
else:
model.load_state_dict(
torch.load(weights_path, map_location=lambda storage, loc: storage))
image = Image.open(IMAGE_NAME)
image = Variable(ToTensor()(image), volatile=True).unsqueeze(0)
if TEST_MODE:
image = image.cuda()
start = time.clock()
out = model(image)
elapsed = (time.clock() - start)
####### Define data loader and models ##########
train_sampler = RandomSampler(
data_source=dataset,
replacement=True,
num_samples=int(1e100), # make the dataloader "infinite"
)
dataloader = DataLoader(
dataset,
batch_size=cfg.batch_size,
sampler=train_sampler,
pin_memory=True,
#shuffle=True,
num_workers=4,
)
G = Generator(
z_dim=cfg.z_dim,
n_feat=cfg.n_feat,
)
if cfg.minibatch_discrimination:
logging.info("Using minibatch discrimination")
D = Discriminator_MD(batch_size=cfg.batch_size, n_feat=cfg.n_feat)
else:
#logging.info("Using minibatch discr")
D = Discriminator(n_feat=cfg.n_feat)
G.apply(weights_init)
D.apply(weights_init)
G = G.to(cfg.device)
D = D.to(cfg.device)
transform = transforms.Compose([
transforms.Resize(resize),
transforms.CenterCrop(resize),
transforms.ToTensor(),
transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
])
imgs = []
for imgfile in args.files:
img = transform(Image.open(imgfile).convert("RGB"))
imgs.append(img)
imgs = torch.stack(imgs, 0).to(device)
g_ema = Generator(args.size, 512, 8)
g_ema.load_state_dict(torch.load(args.ckpt)["g_ema"], strict=False)
g_ema.eval()
g_ema = g_ema.to(device)
with torch.no_grad():
noise_sample = torch.randn(n_mean_latent, 512, device=device)
latent_out = g_ema.style(noise_sample)
latent_mean = latent_out.mean(0)
latent_std = ((latent_out - latent_mean).pow(2).sum() /
n_mean_latent)**0.5
percept = lpips.PerceptualLoss(model="net-lin",
net="vgg",
use_gpu=device.startswith("cuda"))
# custom weights initialization called on netG and netD
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
m.weight.data.normal_(0.0, 0.02)
elif classname.find('BatchNorm') != -1:
m.weight.data.normal_(1.0, 0.02)
m.bias.data.fill_(0)
'''---------------------------------------INSTANTIATE MODELS-------------------------------------'''
from model import Generator, Discriminator
netG = Generator(ngpu).to(device)
netG.apply(weights_init)
if args.netG != '':
netG.load_state_dict(torch.load(args.netG))
# print(netG)
netD = Discriminator(ngpu).to(device)
netD.apply(weights_init)
if args.netD != '':
netD.load_state_dict(torch.load(args.netD))
# print(netD)
if gpu:
netD.to(device)
netG.to(device)
'''----------------------------------------LOSS & OPTIMIZER--------------------------------------'''
