def get_sample_model(config, from_style, to_style, epoch):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
generator_ab = Generator(config.image_size,
config.num_residual_blocks).to(device)
generator_ba = Generator(config.image_size,
config.num_residual_blocks).to(device)
generator_ab_param = glob(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}",
f"generator_ab_{epoch}.pth"))
generator_ba_param = glob(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}",
f"generator_ab_{epoch}.pth"))
print(f"[*] Load checkpoint in {epoch}")
print(f"[*] load generator_ab_{epoch}.pth")
if len(
os.listdir(
os.path.join(config.checkpoint_dir,
f"{from_style}2{to_style}"))) == 0:
raise Exception(f"[!] No checkpoint in {config.checkpoint_dir}")
else:
generator_ab.load_state_dict(
torch.load(generator_ab_param[-1], map_location=device))
generator_ba.load_state_dict(
torch.load(generator_ba_param[-1], map_location=device))
return generator_ab, generator_ba
def build_model(config, from_style, to_style):
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
generator_ab = Generator(config.num_residual_blocks).to(device)
generator_ba = Generator(config.num_residual_blocks).to(device)
discriminator_a = Discriminator(config.image_size).to(device)
discriminator_b = Discriminator(config.image_size).to(device)
generator_ab_param = glob(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}",
f"generator_ab_{config.epoch-1}.pth"))
generator_ba_param = glob(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}",
f"generator_ba_{config.epoch-1}.pth"))
discriminator_a_param = glob(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}",
f"discriminator_a_{config.epoch-1}.pth"))
discriminator_b_param = glob(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}",
f"discriminator_b_{config.epoch-1}.pth"))
print(f"[*] Load checkpoint in {config.checkpoint_dir}")
if not os.path.exists(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}")):
os.makedirs(
os.path.join(config.checkpoint_dir, f"{from_style}2{to_style}"))
if not generator_ab_param:
print(f"[!] No checkpoint in {config.checkpoint_dir}")
generator_ab.apply(weights_init)
generator_ba.apply(weights_init)
discriminator_a.apply(weights_init)
discriminator_b.apply(weights_init)
else:
generator_ab.load_state_dict(
torch.load(generator_ab_param[-1], map_location=device))
generator_ba.load_state_dict(
torch.load(generator_ba_param[-1], map_location=device))
discriminator_a.load_state_dict(
torch.load(discriminator_a_param[-1], map_location=device))
discriminator_b.load_state_dict(
torch.load(discriminator_b_param[-1], map_location=device))
return generator_ab, generator_ba, discriminator_a, discriminator_b
def test_all_level_no_mask_yes_attr(args):
"""Test model with input image and attributes."""
transform = transforms.Compose(
[Normalize(0.5, 0.5),
CenterSquareMask(),
ScaleNRotate(),
ToTensor()])
batch_size = 1
num_attrs = 40
resolutions_to = [4, 8, 8, 16, 16, 32, 32, 64, 64, 128, 128, 256,
256] # 512, 512]
levels = [1, 1.5, 2, 2.5, 3, 3.5, 4, 4.5, 5, 5.5, 6, 6.5, 7] # 7.5, 8]
data_shape = [batch_size, 3, 512, 512]
G = Generator(data_shape, use_mask=False, num_attrs=num_attrs)
D = Discriminator(data_shape, num_attrs=num_attrs)
for res, lev in zip(resolutions_to, levels):
dataset = CelebAHQDataset(args.data_dir, res, transform)
dataloader = DataLoader(dataset, batch_size, True)
sample = iter(dataloader).next() # noqa: B305
image = sample['image']
masked_image = sample['masked_image']
mask = sample['mask']
attr = sample['attr']
print(f"level: {lev}, resolution: {res}, image: {masked_image.shape}, \
mask: {mask.shape}")
# Generator
if isinstance(lev, int):
# training state
fake_image1 = G(masked_image, attr, cur_level=lev)
assert list(fake_image1.shape) == [batch_size, 3, res, res], \
f'{res, lev} test failed'
else:
# transition state
fake_image2 = G(masked_image, attr, cur_level=lev)
assert list(fake_image2.shape) == [batch_size, 3, res, res], \
f'{res, lev} test failed'
# Discriminator
if isinstance(lev, int):
# training state
cls1, attr1 = D(image, lev)
assert list(cls1.shape) == [batch_size, 1], \
f'{res, lev} test failed'
assert list(attr1.shape) == [batch_size, num_attrs], \
f'{res, lev} test failed'
else:
# transition state
cls2, attr2 = D(image, lev)
assert list(cls2.shape) == [batch_size, 1], \
f'{res, lev} test failed'
assert list(attr2.shape) == [batch_size, num_attrs], \
f'{res, lev} test failed'
def synthesize(args):
hp.batch_size = 1
model = Generator(hp)
model.training = False
model.load_parameters(args.f_model, raise_if_missing=True)
wave = np.load(args.f_mel)
x_mel = nn.Variable.from_numpy_array(wave[None, ...])
o_aud = model(x_mel)
o_aud.forward(clear_buffer=True)
wavfile.write(args.f_output, rate=hp.sr, data=o_aud.d[0, 0].copy())
def load_generator(args, config):
if not os.path.exists(args.weights):
logger.info(f"Generator weights {args.weights} does not exist!")
return None
generator = Generator(config)
generator_state_dict = torch.load(args.weights, map_location=args.device)
generator.load_state_dict(generator_state_dict)
generator.to(args.device)
return generator
def test_end_to_end(args):
"""Test end to end data handling process."""
batch_size = 1
resolutions = [256, 256]
levels = [7, 7]
num_classes = 5
num_layers = 1
data_shape = [batch_size, 3, 256, 256]
transform = transforms.Compose([
Normalize(0.5, 0.5),
TargetMask(num_classes),
ScaleNRotate(),
ToTensor()
])
G = Generator(data_shape)
D = Discriminator(data_shape, num_classes, num_layers)
for res, lev in zip(resolutions, levels):
dataset = VGGFace2Dataset(args.data_dir, res, args.landmark_info_path,
args.identity_info_path, transform)
dataloader = DataLoader(dataset, batch_size, True)
sample = iter(dataloader).next() # noqa: B305
image = sample['image']
real_mask = sample['real_mask']
obs_mask = sample['obs_mask']
target_id = sample['target_id']
print(f"lev: {lev}, res: {res}, image: {image.shape}, \
mask: {real_mask.shape}, {obs_mask.shape}, \
target_id: {target_id}")
# Generator
fake_image = G(image, obs_mask, cur_level=lev)
assert list(fake_image.shape) == [batch_size, 3, res, res], \
f'Generator: {res, lev} test failed'
# Discriminator (original)
cls1, pix_cls1 = D(image, lev)
assert list(cls1.shape) == [batch_size, 1], \
f'Discriminator: {res, lev} test failed'
assert list(pix_cls1.shape) == [batch_size, num_classes, res, res], \
f'Pixel-Discriminator: {res, lev} test failed'
cls2, pix_cls2 = D(fake_image, lev)
assert list(cls2.shape) == [batch_size, 1], \
f'Discriminator: {res, lev} test failed'
assert list(pix_cls2.shape) == [batch_size, num_classes, res, res], \
f'Pixel-Discriminator: {res, lev} test failed'
def build_model(self):
self.G = Generator()
self.D = Discriminator()
self.C = DomainClassifier()
self.g_optimizer = flow.optim.Adam(self.G.parameters(), self.g_lr,
[self.beta1, self.beta2])
self.d_optimizer = flow.optim.Adam(self.D.parameters(), self.d_lr,
[self.beta1, self.beta2])
self.c_optimizer = flow.optim.Adam(self.C.parameters(), self.c_lr,
[self.beta1, self.beta2])
self.print_network(self.G, "G")
self.print_network(self.D, "D")
self.print_network(self.C, "C")
self.G.to(self.device)
self.D.to(self.device)
self.C.to(self.device)
if form["eyesColor"] != "random":
utag.append(form["eyesColor"])
# print("utag:", utag)
generate_image("static/webout.png", 10, 10, utag)
return redirect(
url_for('home',
hairDefault=form["hairColor"],
eyesDefault=form["eyesColor"]))
return render_template('home.html',
imgt=str(time.time()),
hairColors=utils.hair,
eyesColors=utils.eyes,
hairDefault=hairDefault,
eyesDefault=eyesDefault)
if __name__ == '__main__':
if len(sys.argv) == 2:
netGpath = sys.argv[1]
NETG = Generator(NZ, len(utils.hair) + len(utils.eyes)).to(device)
try:
NETG.load_state_dict(
torch.load(netGpath,
map_location=lambda storage, loc: storage))
app.run(debug=True)
except:
print("`{}` not found".format(netGpath))
else:
print("Usage: python3 webapp.py [netG model path]")
train_set = TrainDatasetFromFolder('data/DIV2K_train_HR',
crop_size=CROP_SIZE,
upscale_factor=UPSCALE_FACTOR)
val_set = ValDatasetFromFolder('data/DIV2K_valid_HR',
crop_size=CROP_SIZE,
upscale_factor=UPSCALE_FACTOR)
train_loader = DataLoader(dataset=train_set,
num_workers=4,
batch_size=4,
shuffle=True)
val_loader = DataLoader(dataset=val_set,
num_workers=4,
batch_size=1,
shuffle=False)
netG = Generator(num_rrdb_blocks=16, scaling_factor=UPSCALE_FACTOR)
print('# generator parameters:',
sum(param.numel() for param in netG.parameters()))
netD = Discriminator(opt.crop_size)
print('# discriminator parameters:',
sum(param.numel() for param in netD.parameters()))
netD_HF = Discriminator(image_size=dwt_size)
generator_criterion = EGeneratorLoss()
discriminator_criterion = torch.nn.BCEWithLogitsLoss()
if torch.cuda.is_available():
netG.cuda()
netD.cuda()
netD_HF.cuda()
generator_criterion.cuda()
train_set = TrainDatasetFromFolder('../../../CelebA-HQ-img/',
crop_size=CROP_SIZE,
upscale_factor=UPSCALE_FACTOR)
val_set = ValDatasetFromFolder('../../../CelebA-HQ-img/',
crop_size=CROP_SIZE,
upscale_factor=UPSCALE_FACTOR)
train_loader = DataLoader(dataset=train_set,
num_workers=4,
batch_size=opt.batchSize,
shuffle=True)
val_loader = DataLoader(dataset=val_set,
num_workers=4,
batch_size=opt.testBatchSize,
shuffle=False)
netG = Generator(UPSCALE_FACTOR).to(device)
netD = Discriminator().to(device)
optimizerG = optim.RMSprop(netG.parameters(), lr=opt.lr)
criterion = nn.MSELoss()
# loop over the dataset multiple times
for epoch in range(opt.start, NUM_EPOCHS + 1):
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
# zero the parameter gradients
optimizerG.zero_grad()
print("Train set size: " + str(len(trainset)))
# Whether training form checkpoint
if opt.resume:
# Load checkpoint.
print('==> Resuming from checkpoint..')
assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
checkpoint = torch.load('./checkpoint/ckpt.t7')
netD = checkpoint['netD']
netG = checkpoint['netG']
start_epoch = checkpoint['epoch']
else:
print('==> Building model..')
# Create an instance of the nn.module class defined above:
netD = Discriminator(trainset, opt.batchSize, reuse=isTrain)
netG = Generator(randomInput, opt.batchSize, reuse=True)
start_epoch = 0
# For training on GPU, we need to transfer net and data onto the GPU
# http://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#training-on-gpu
if opt.ngpu >= 1:
netD = netD.cuda()
netD = torch.nn.DataParallel(netD,
device_ids=range(torch.cuda.device_count()))
netG = netG.cuda()
netG = torch.nn.DataParallel(netG,
device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
# Optimizers
G_optimizer = torch.optim.RMSprop(netG.parameters(), lr=opt.lr)
help='using GPU or CPU')
parser.add_argument('--image_name',
type=str,
help='test low resolution image name')
parser.add_argument('--model_name',
default='FAN_PSNR_X4.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
MODEL_NAME = opt.model_name
model = Generator(UPSCALE_FACTOR).eval()
if TEST_MODE:
model.cuda()
model.load_state_dict(torch.load('epochs/' + MODEL_NAME))
else:
model.load_state_dict(
torch.load('epochs/' + MODEL_NAME,
map_location=lambda storage, loc: storage))
image = Image.open(IMAGE_NAME)
image = image.resize((512, 512))
origin = image
origin = Variable(ToTensor()(origin), volatile=True).unsqueeze(0)
x, y = image.size
image = image.resize((x // opt.upscale_factor, y // opt.upscale_factor),
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--checkpoint_dir", type=str, default="output/ckpt")
parser.add_argument("--model_config",
type=str,
default="model_config.json")
parser.add_argument("--no_cuda",
action='store_true',
help="Avoid using CUDA when available")
parser.add_argument('--photo_dir',
type=str,
default="data/photo",
help='path to photo datasets.')
parser.add_argument('--edge_smooth_dir',
type=str,
default="data/edge_smooth",
help='path to edge_smooth datasets.')
parser.add_argument('--target_dir',
type=str,
default="data/target",
help='path to target datasets.')
parser.add_argument('--content_loss_weight',
type=float,
default=10,
help='content loss weight')
parser.add_argument('--seed', type=int, default=42, help='seed')
parser.add_argument('--adam_beta',
type=float,
default=0.5,
help='adam_beta')
parser.add_argument('--n_epochs',
type=int,
default=100,
help='number of epochs of training')
parser.add_argument('--n_init_epoch',
type=int,
default=15,
help='number of epochs of initializing')
parser.add_argument('--batch_size',
type=int,
default=8,
help='size of the batches')
parser.add_argument('--lr',
type=float,
default=0.0002,
help='initial learning rate')
parser.add_argument(
'--n_cpu',
type=int,
default=0,
help='number of cpu threads to use during batch generation')
parser.add_argument('--logging_steps',
type=int,
default=50,
help="Log every X updates steps.")
parser.add_argument('--save_steps',
type=int,
default=3000,
help='Save checkpoint every X updates steps.')
args = parser.parse_args()
# Setup logging
logging.basicConfig(
format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
datefmt='%m/%d/%Y %H:%M:%S',
level=logging.INFO)
model_config = Config.load(args.model_config)
args.device = torch.device(
"cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
args.n_gpu = torch.cuda.device_count()
set_seed(args)
logger.warning("device: %s, n_gpu: %s", args.device, args.n_gpu)
generator = Generator(model_config).to(args.device)
discriminator = Discriminator(model_config).to(args.device)
feature_extractor = FeatureExtractor(model_config).to(args.device)
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
photo_dataloader, _ = load_image_dataloader(args.photo_dir, transform,
args.batch_size, args.n_cpu)
edge_smooth_dataloader, _ = load_image_dataloader(args.edge_smooth_dir,
transform,
args.batch_size,
args.n_cpu)
animation_dataloader, _ = load_image_dataloader(args.target_dir, transform,
args.batch_size,
args.n_cpu)
train(args, generator, discriminator, feature_extractor, photo_dataloader,
edge_smooth_dataloader, animation_dataloader, args.checkpoint_dir)
def __init__(self, config):
"""Class initializer.
1. Read self.configurations from self.config.py
2. Check gpu availability
3. Create a model and training related objects
- Model (Generator, Discriminator)
- Optimizer
- Loss and loss histories
- Replay memory
- Snapshot
"""
self.config = config
self.D_repeats = self.config.train.D_repeats
self.total_size = int(self.config.train.total_size *
self.config.train.dataset_unit)
self.train_size = int(self.config.train.train_size *
self.config.train.dataset_unit)
self.transition_size = int(self.config.train.transition_size *
self.config.train.dataset_unit)
assert (self.total_size == (self.train_size + self.transition_size)) \
and self.train_size > 0 and self.transition_size > 0
# GPU
self.check_gpu()
self.mode = self.config.train.mode
self.use_mask = self.config.train.use_mask
# Data Shape
dataset_shape = [
1, self.config.dataset.num_channels,
self.config.train.net.max_resolution,
self.config.train.net.max_resolution
]
# Generator & Discriminator Creation
self.G = Generator(dataset_shape,
fmap_base=self.config.train.net.fmap_base,
fmap_min=self.config.train.net.min_resolution,
fmap_max=self.config.train.net.max_resolution,
latent_size=self.config.train.net.latent_size,
use_mask=self.use_mask,
leaky_relu=True,
instancenorm=True)
spectralnorm = True if self.config.loss.gan == Gan.sngan else False
self.D = Discriminator(dataset_shape,
num_classes=self.config.dataset.num_classes,
num_layers=self.config.train.net.num_layers,
fmap_base=self.config.train.net.fmap_base,
fmap_min=self.config.train.net.min_resolution,
fmap_max=self.config.train.net.max_resolution,
latent_size=self.config.train.net.latent_size,
leaky_relu=True,
instancenorm=True,
spectralnorm=spectralnorm)
self.register_on_gpu()
self.create_optimizer()
# Loss
self.loss = FaceGenLoss(self.config, self.use_cuda,
self.config.env.num_gpus)
# Replay Memory
self.replay_memory = ReplayMemory(self.config, self.use_cuda,
self.config.replay.enabled)
self.global_it = 1
self.global_cur_nimg = 1
# restore
self.snapshot = Snapshot(self.config, self.use_cuda)
self.snapshot.prepare_logging()
self.snapshot.restore_model(self.G, self.D, self.optim_G, self.optim_D)
def __init__(self, args):
"""
Args:
args (Namespace): Program arguments from argparser
"""
# Store args
self.num_epochs = args.num_epochs
self.start_epoch = args.start_epoch
self.generator_lr = args.generator_lr
self.discriminator_lr = args.discriminator_lr
self.decay_after = args.decay_after
self.mini_batch_size = args.batch_size
self.cycle_loss_lambda = args.cycle_loss_lambda
self.identity_loss_lambda = args.identity_loss_lambda
self.device = args.device
self.epochs_per_save = args.epochs_per_save
self.sample_rate = args.sample_rate
self.validation_A_dir = os.path.join(args.origin_data_dir,
args.speaker_A_id)
self.output_A_dir = os.path.join(args.output_data_dir,
args.speaker_A_id)
self.validation_B_dir = os.path.join(args.origin_data_dir,
args.speaker_B_id)
self.output_B_dir = os.path.join(args.output_data_dir,
args.speaker_B_id)
self.infer_data_dir = args.infer_data_dir
self.pretrain_models = args.pretrain_models
# Initialize speakerA's dataset
self.dataset_A = self.loadPickleFile(
os.path.join(
args.preprocessed_data_dir,
args.speaker_A_id,
f"{args.speaker_A_id}_normalized.pickle",
))
dataset_A_norm_stats = np.load(
os.path.join(
args.preprocessed_data_dir,
args.speaker_A_id,
f"{args.speaker_A_id}_norm_stat.npz",
))
self.dataset_A_mean = dataset_A_norm_stats["mean"]
self.dataset_A_std = dataset_A_norm_stats["std"]
# Initialize speakerB's dataset
self.dataset_B = self.loadPickleFile(
os.path.join(
args.preprocessed_data_dir,
args.speaker_B_id,
f"{args.speaker_B_id}_normalized.pickle",
))
dataset_B_norm_stats = np.load(
os.path.join(
args.preprocessed_data_dir,
args.speaker_B_id,
f"{args.speaker_B_id}_norm_stat.npz",
))
self.dataset_B_mean = dataset_B_norm_stats["mean"]
self.dataset_B_std = dataset_B_norm_stats["std"]
# Compute lr decay rate
self.n_samples = len(self.dataset_A)
print(f"n_samples = {self.n_samples}")
self.generator_lr_decay = self.generator_lr / float(
self.num_epochs * (self.n_samples // self.mini_batch_size))
self.discriminator_lr_decay = self.discriminator_lr / float(
self.num_epochs * (self.n_samples // self.mini_batch_size))
print(f"generator_lr_decay = {self.generator_lr_decay}")
print(f"discriminator_lr_decay = {self.discriminator_lr_decay}")
# Initialize Train Dataloader
self.num_frames = args.num_frames
self.dataset = VCDataset(
datasetA=self.dataset_A,
datasetB=self.dataset_B,
n_frames=args.num_frames,
max_mask_len=args.max_mask_len,
)
self.train_dataloader = flow.utils.data.DataLoader(
dataset=self.dataset,
batch_size=self.mini_batch_size,
shuffle=True,
drop_last=False,
)
# Initialize Generators and Discriminators
self.generator_A2B = Generator().to(self.device)
self.generator_B2A = Generator().to(self.device)
self.discriminator_A = Discriminator().to(self.device)
self.discriminator_B = Discriminator().to(self.device)
# Discriminator to compute 2 step adversarial loss
self.discriminator_A2 = Discriminator().to(self.device)
# Discriminator to compute 2 step adversarial loss
self.discriminator_B2 = Discriminator().to(self.device)
# Initialize Optimizers
g_params = list(self.generator_A2B.parameters()) + list(
self.generator_B2A.parameters())
d_params = (list(self.discriminator_A.parameters()) +
list(self.discriminator_B.parameters()) +
list(self.discriminator_A2.parameters()) +
list(self.discriminator_B2.parameters()))
self.generator_optimizer = flow.optim.Adam(g_params,
lr=self.generator_lr,
betas=(0.5, 0.999))
self.discriminator_optimizer = flow.optim.Adam(
d_params, lr=self.discriminator_lr, betas=(0.5, 0.999))
default=1,
help='number of GPUs to use')
parser.add_argument('--netG', default=None, help="path to netG")
opt = parser.parse_args()
randomSeed = random.randint(1, 10000)
print("Random Seed: ", randomSeed)
random.seed(randomSeed)
torch.manual_seed(randomSeed)
if opt.netG is not None:
nz = opt.nz
ngf = opt.ngf
nc = 3
# netG = dcgan.Generator(opt.ngpu, nz, nc, ngf)
netG = Generator(nz, len(utils.hair) + len(utils.eyes)).to(device)
netG.load_state_dict(
torch.load(opt.netG, map_location=lambda storage, loc: storage))
def random_sample(sample_number):
samples = []
for _ in range(0, sample_number):
path = os.path.join(opt.dataroot, "")
files = [name for name in os.listdir(path)]
index = random.randint(0, len(files) - 1)
samples.append(os.path.join(path, files[index]))
return samples
def main(writer):
dataset = AnimeDataset(avatar_tag_dat_path,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5),
(0.5, 0.5, 0.5))
]))
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers,
drop_last=True)
G = Generator(noise_size, len(utils.hair) + len(utils.eyes)).to(device)
D = Discriminator(len(utils.hair), len(utils.eyes)).to(device)
G_optim = torch.optim.Adam(G.parameters(),
lr=learning_rate_g,
betas=(beta_1, 0.999))
D_optim = torch.optim.Adam(D.parameters(),
lr=learning_rate_d,
betas=(beta_1, 0.999))
criterion = nn.BCELoss()
# training
iteration = 0
real_label = torch.ones(batch_size).to(device)
# real_label = torch.Tensor(batch_size).uniform_(0.9, 1).to(device) # soft labeling
fake_label = torch.zeros(batch_size).to(device)
for epoch in range(max_epoch + 1):
for i, (real_tag, real_img) in enumerate(data_loader):
real_img = real_img.to(device)
real_tag = real_tag.to(device)
# train D with real images
D.zero_grad()
real_score, real_predict = D(real_img)
real_discrim_loss = criterion(real_score, real_label)
real_classifier_loss = criterion(real_predict, real_tag)
# train D with fake images
z, fake_tag = utils.fake_generator(batch_size, noise_size, device)
fake_img = G(z, fake_tag).to(device)
fake_score, fake_predict = D(fake_img)
fake_discrim_loss = criterion(fake_score, fake_label)
discrim_loss = (real_discrim_loss + fake_discrim_loss) * 0.5
classifier_loss = real_classifier_loss * lambda_cls
# gradient penalty
alpha_size = [1] * real_img.dim()
alpha_size[0] = real_img.size(0)
alpha = torch.rand(alpha_size).to(device)
x_hat = Variable(alpha * real_img.data + (1 - alpha) *
(real_img.data + 0.5 * real_img.data.std() *
torch.rand(real_img.size()).to(device)),
requires_grad=True).to(device)
fake_score, fake_tag = D(x_hat)
gradients = grad(outputs=fake_score,
inputs=x_hat,
grad_outputs=torch.ones(
fake_score.size()).to(device),
create_graph=True,
retain_graph=True,
only_inputs=True)[0].view(x_hat.size(0), -1)
gradient_penalty = lambda_gp * (
(gradients.norm(2, dim=1) - 1)**2).mean()
D_loss = discrim_loss + classifier_loss + gradient_penalty
D_loss.backward()
D_optim.step()
# train G
G.zero_grad()
z, fake_tag = utils.fake_generator(batch_size, noise_size, device)
fake_img = G(z, fake_tag).to(device)
fake_score, fake_predict = D(fake_img)
discrim_loss = criterion(fake_score, real_label)
classifier_loss = criterion(fake_predict, fake_tag) * lambda_cls
G_loss = discrim_loss + classifier_loss
G_loss.backward()
G_optim.step()
# plot loss curve
writer.add_scalar('Loss_D', D_loss.item(), iteration)
writer.add_scalar('Loss_G', G_loss.item(), iteration)
print('[{}/{}][{}/{}] Iteration: {}'.format(
epoch, max_epoch, i, len(data_loader), iteration))
if iteration % interval == interval - 1:
fake_img = G(fix_noise, fix_tag)
vutils.save_image(utils.denorm(fake_img[:64, :, :, :]),
os.path.join(
image_path,
'fake_image_{}.png'.format(iteration)),
padding=0)
vutils.save_image(utils.denorm(real_img[:64, :, :, :]),
os.path.join(
image_path,
'real_image_{}.png'.format(iteration)),
padding=0)
grid = vutils.make_grid(utils.denorm(fake_img[:64, :, :, :]),
padding=0)
writer.add_image('generation results', grid, iteration)
iteration += 1
# checkpoint
torch.save(G.state_dict(),
os.path.join(model_path, 'netG_epoch_{}.pth'.format(epoch)))
torch.save(D.state_dict(),
os.path.join(model_path, 'netD_epoch_{}.pth'.format(epoch)))
landmark_info_path = './dataset/VGGFACE2/train/all_loose_landmarks_256.csv'
identity_info_path = './dataset/VGGFACE2/identity_info.csv'
filtered_list = './dataset/VGGFACE2/train/all_filtered_results.csv'
transform = transforms.Compose([Normalize(0.5, 0.5), ToTensor()])
batch_size = 2
num_classes = 2
num_attrs = 1
resolutions_to = [4, 8, 8, 16, 16, 32, 32]
levels = [1, 1.125, 2, 2.5, 3, 3.5, 4]
data_shape = [batch_size, 3, 32, 32]
G = Generator(data_shape,
use_mask=False,
use_attrs=True,
num_attrs=num_attrs,
latent_size=256)
D = Discriminator(data_shape,
use_attrs=True,
num_attrs=num_attrs,
latent_size=256)
for res, lev in zip(resolutions_to, levels):
dataset = VGGFace2Dataset('./dataset/VGGFACE2/train',
res,
landmark_info_path,
identity_info_path,
filtered_list,
transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
def __init__(
self,
logf0s_normalization,
mcep_normalization,
coded_sps_A_norm,
coded_sps_B_norm,
model_checkpoint,
validation_A_dir,
output_A_dir,
validation_B_dir,
output_B_dir,
restart_training_at=None,
):
self.start_epoch = 0
self.num_epochs = 200000
self.mini_batch_size = 10
self.dataset_A = self.loadPickleFile(coded_sps_A_norm)
self.dataset_B = self.loadPickleFile(coded_sps_B_norm)
self.device = flow.device(
"cuda" if flow.cuda.is_available() else "cpu")
# Speech Parameters
logf0s_normalization = np.load(logf0s_normalization)
self.log_f0s_mean_A = logf0s_normalization["mean_A"]
self.log_f0s_std_A = logf0s_normalization["std_A"]
self.log_f0s_mean_B = logf0s_normalization["mean_B"]
self.log_f0s_std_B = logf0s_normalization["std_B"]
mcep_normalization = np.load(mcep_normalization)
self.coded_sps_A_mean = mcep_normalization["mean_A"]
self.coded_sps_A_std = mcep_normalization["std_A"]
self.coded_sps_B_mean = mcep_normalization["mean_B"]
self.coded_sps_B_std = mcep_normalization["std_B"]
# Generator and Discriminator
self.generator_A2B = Generator().to(self.device)
self.generator_B2A = Generator().to(self.device)
self.discriminator_A = Discriminator().to(self.device)
self.discriminator_B = Discriminator().to(self.device)
# Loss Functions
criterion_mse = flow.nn.MSELoss()
# Optimizer
g_params = list(self.generator_A2B.parameters()) + list(
self.generator_B2A.parameters())
d_params = list(self.discriminator_A.parameters()) + list(
self.discriminator_B.parameters())
# Initial learning rates
self.generator_lr = 2e-4
self.discriminator_lr = 1e-4
# Learning rate decay
self.generator_lr_decay = self.generator_lr / 200000
self.discriminator_lr_decay = self.discriminator_lr / 200000
# Starts learning rate decay from after this many iterations have passed
self.start_decay = 10000
self.generator_optimizer = flow.optim.Adam(g_params,
lr=self.generator_lr,
betas=(0.5, 0.999))
self.discriminator_optimizer = flow.optim.Adam(
d_params, lr=self.discriminator_lr, betas=(0.5, 0.999))
# To Load save previously saved models
self.modelCheckpoint = model_checkpoint
os.makedirs(self.modelCheckpoint, exist_ok=True)
# Validation set Parameters
self.validation_A_dir = validation_A_dir
self.output_A_dir = output_A_dir
os.makedirs(self.output_A_dir, exist_ok=True)
self.validation_B_dir = validation_B_dir
self.output_B_dir = output_B_dir
os.makedirs(self.output_B_dir, exist_ok=True)
# Storing Discriminatior and Generator Loss
self.generator_loss_store = []
self.discriminator_loss_store = []
self.file_name = "log_store_non_sigmoid.txt"
from torchvision.utils import save_image
import torchvision.transforms.functional as TF
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--param_path', default='parameters/monet.pth')
parser.add_argument('--input_dir')
parser.add_argument('--output_dir', default='results')
args = parser.parse_args()
images_name = sorted(os.listdir(args.input_dir))
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
with torch.no_grad():
net = Generator().to(device).eval()
net.load_state_dict(torch.load(args.param_path))
for image_name in images_name:
image = Image.open(os.path.join(args.input_dir,
image_name)).convert('RGB')
image = TF.to_tensor(image).to(device).unsqueeze(dim=0)
image = net(image)
save_image(image, os.path.join(args.output_dir, image_name))
print(f'save {image_name}')
