代码：https://github.com/alisadeghian/PGMGAN
论文
论文解说
Partition-Guided GANs
readme
基于GANS的无监督多图像生成：分区引导混合的生成对抗网络
（Partition Guided Mixture of Generative Adversarial Networks）
我们提出的完全无监督的图像生成模型PGMGAN学习了基于语义相似性的空间划分，并从每个分区生成图像，以减少模式崩溃和模式连接。我们提出了一种新的分块/引导方法，保证为生成器提供指向其指定区域的方向 region. scan_guide_biggan。
Getting Started
Installation
Clone this repo:
git clone https://github.com/alisadeghian/PGMGAN.git
cd PGMGAN
Install the dependencies
conda create --name PGMGAN python=3.7
conda activate PGMGAN
conda install --file requirements.txt
conda install -c conda-forge tensorboardx
Training and Evaluation
Train a model on CIFAR:
python train.py configs/cifar/scan_guide_biggan.yaml
Visualize samples and inferred 推理的 clusters:
python visualize_clusters.py configs/cifar/scan_guide_biggan.yaml --show_clusters
The samples and clusters will be saved to output/cifar/scan_guide_biggan/clusters.
Evaluate the model’s FID: You will need to first gather a set of ground truth train set images to compute metrics against.
您将需要首先收集一组地面真相训练集图像来计算度量标准。
python utils/get_gt_imgs.py --cifar
Then, run the evaluation script:
python metrics.py configs/cifar/scan_guide_biggan.yaml --fid --every -1
You can also evaluate with other metrics 度量 by appending additional flags, such as Inception Score (--inception), the number of covered modes + reverse-KL divergence (--modes), and cluster metrics (--cluster_metrics).
Pretrained Models
You can download pretrained models on CIFAR from here [https://drive.google.com/drive/folders/1lsZKU6T0H91ThW_lCpsEJNGg0rTYGBcY?usp=sharing]
and place them in the output/cifar/scan_guide_biggan/chkpts/ directory.
To reproduce the results in the paper use the following command:
python metrics.py configs/cifar/scan_guide_biggan.yaml --fid --every -1
Evaluation
Visualizations
To visualize generated samples and inferred clusters, run
python visualize_clusters.py config-file
You can set the flag --show_clusters to also visualize the real inferred clusters, but this requires that you have a path to training set images.
Metrics
To obtain generation metrics, fill in the path to your ImageNet or Places dataset directories in utils/get_gt_imgs.py and then run
python utils/get_gt_imgs.py --imagenet --places
to precompute batches of GT images for FID/FSD evaluation.
Then, you can use
python metrics.py config-file
with the appropriate flags compute the FID (--fid), FSD (--fsd), IS (--inception), number of modes covered/ reverse-KL divergence (--modes) and clustering metrics (–cluster_metrics) for each of the checkpoints.
Acknowledgments
This code is heavily based on the GAN-stability and self-cond-gan code bases. Our FSD code is taken from the GANseeing work. To compute inception score, we use the code provided from Shichang Tang. To compute FID, we use the code provided from TTUR. We also use pretrained classifiers given by the pytorch-playground.
补充知识
pprint.PrettyPrinter(indent=1)
美化输出
控制输出的打印时的缩进，行宽、打印深度等，
定义类如下
class pprint.PrettyPrinter（indent = 1，width = 80，depth = None，stream = None，*，compact = False ）
- indent 缩进
- width 宽度
- depth 打印深度
- stream 指输出流对象，stream = None输出流对象默认是sys.stdout
- compact
如果compact为false（默认值），则长序列中的每个项目将在单独的行上进行格式化。如果compact为true，则将在每个输出行上格式化适合宽度的项目。
PGMGAN-main
错误
训练
python train.py configs/cifar/scan_guide_biggan.yaml
评估
python metrics.py configs/cifar/scan_guide_biggan.yaml --fid --every -1
需要先生成 gt 图像
python utils/get_gt_imgs.py --cifar
train.py
import argparse
import os
import copy
import pprint
from os import path
import torch
from torch import nn
from gan_training import utils
from gan_training.train import Trainer, update_average
from gan_training.logger import Logger
from gan_training.checkpoints import CheckpointIO
from gan_training.inputs import get_dataset
from gan_training.distributions import get_ydist, get_zdist
from gan_training.eval import Evaluator
from gan_training.config import (load_config, get_clusterer, build_models, build_optimizers)
from seeing.pidfile import exit_if_job_done, mark_job_done
torch.backends.cudnn.benchmark = True
# Arguments
parser = argparse.ArgumentParser(
description='Train a GAN with different regularization strategies.')#正则化策略
parser.add_argument('config', type=str, help='Path to config file.')
parser.add_argument('--outdir', type=str, help='used to override outdir (useful for multiple runs)')
parser.add_argument('--nepochs', type=int, default=3000, help='number of epochs to run before terminating')#结束
parser.add_argument('--model_it', type=int, default=-1,
help='which model iteration to load from, -1 loads the most recent model')
#从哪个模型迭代加载，-1加载最近的模型
parser.add_argument('--devices', nargs='+', type=str, default=['0'], help='devices to use')
args = parser.parse_args()
config = load_config(args.config, 'configs/default.yaml')
out_dir = config['training']['out_dir'] if args.outdir is None else args.outdir
def main():
pp = pprint.PrettyPrinter(indent=1)
pp.pprint({
'data': config['data'],
'generator': config['generator'],
'discriminator': config['discriminator'],
'clusterer': config['clusterer'],
'training': config['training']
})
is_cuda = torch.cuda.is_available()
# Short hands
batch_size = config['training']['batch_size']
log_every = config['training']['log_every']
inception_every = config['training']['inception_every']
backup_every = config['training']['backup_every']
sample_nlabels = config['training']['sample_nlabels']
nlabels = config['data']['nlabels']
sample_nlabels = min(nlabels, sample_nlabels)
weight_guide_init = config['resnet']['weight_guide']
checkpoint_dir = path.join(out_dir, 'chkpts')
# Create missing directories
if not path.exists(out_dir):
os.makedirs(out_dir)
if not path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
# Logger
checkpoint_io = CheckpointIO(checkpoint_dir=checkpoint_dir)
device = torch.device("cuda:0" if is_cuda else "cpu")
train_dataset, _ = get_dataset(
name=config['data']['type'],
data_dir=config['data']['train_dir'],
size=config['data']['img_size'],
deterministic=config['data']['deterministic'])
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=batch_size,
num_workers=config['training']['nworkers'],
shuffle=True,
pin_memory=True,
sampler=None,
drop_last=True)
# Create models
generator, discriminator = build_models(config)
# Put models on gpu if needed
generator = generator.to(device)
discriminator = discriminator.to(device)
for name, module in discriminator.named_modules():
if isinstance(module, nn.Sigmoid):
print('Found sigmoid layer in discriminator; not compatible with BCE with logits')
exit()
g_optimizer, d_optimizer = build_optimizers(generator, discriminator, config)
devices = [int(x) for x in args.devices]
generator = nn.DataParallel(generator, device_ids=devices)
discriminator = nn.DataParallel(discriminator, device_ids=devices)
# Register modules to checkpoint
checkpoint_io.register_modules(generator=generator,
discriminator=discriminator,
g_optimizer=g_optimizer,
d_optimizer=d_optimizer)
# Logger
logger = Logger(log_dir=path.join(out_dir, 'logs'),
img_dir=path.join(out_dir, 'imgs'),
monitoring=config['training']['monitoring'],
monitoring_dir=path.join(out_dir, 'monitoring'))
# Distributions
ydist = get_ydist(nlabels, device=device)
zdist = get_zdist(config['z_dist']['type'], config['z_dist']['dim'],
device=device)
ntest = config['training']['ntest']
x_test, y_test = utils.get_nsamples(train_loader, ntest)
x_test, y_test = x_test.to(device), y_test.to(device)
z_test = zdist.sample((ntest, ))
utils.save_images(x_test, path.join(out_dir, 'real.png'))
logger.add_imgs(x_test, 'gt', 0)
# Test generator
if config['training']['take_model_average']:
print('Taking model average')
bad_modules = [nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d]
for model in [generator, discriminator]:
for name, module in model.named_modules():
for bad_module in bad_modules:
if isinstance(module, bad_module):
print('Batch norm in discriminator not compatible with exponential moving average')
exit()
generator_test = copy.deepcopy(generator)
checkpoint_io.register_modules(generator_test=generator_test)
else:
generator_test = generator
# Load partitioner.
clusterer = get_clusterer(config)(**config['clusterer']['kwargs'])
clusterer = clusterer.cuda()
clusterer.scan_selflabel_model.eval()
clusterer.fill_x_labels(train_loader)
# Test the accuracy of the pre-trained partitioner.
clusterer.test_acc(train_loader)
# Load checkpoint if it exists
it = utils.get_most_recent(checkpoint_dir, 'model') if args.model_it == -1 else args.model_it
epoch_idx = -1
# Evaluator
evaluator = Evaluator(
generator_test,
zdist,
ydist,
train_loader=train_loader,
clusterer=clusterer,
batch_size=batch_size,
device=device,
inception_nsamples=config['training']['inception_nsamples'])
# Trainer
trainer = Trainer(generator,
discriminator,
g_optimizer,
d_optimizer,
gan_type=config['training']['gan_type'],
reg_type=config['training']['reg_type'],
reg_param=config['training']['reg_param'])
# Training loop
print('Start training...')
while it < args.nepochs * len(train_loader):
epoch_idx += 1
for x_real, y in train_loader:
it += 1
x_real, y = x_real.to(device), y.to(device)
z = zdist.sample((batch_size, ))
y = clusterer.get_labels(x_real, y).to(device)
# Discriminator updates
dloss, reg = trainer.discriminator_trainstep(x_real, y, z)
logger.add('losses', 'discriminator', dloss, it=it)
logger.add('losses', 'regularizer', reg, it=it)
if it % 4 == 0:
# Generators updates
weight_guide = max(
config['resnet']['min_weight_guide'],
weight_guide_init + config['resnet']['annealing'] * it/4)
gloss, guide_loss = trainer.generator_trainstep_guided(
y, z, clusterer, weight_guide)
logger.add('losses', 'guide_loss', guide_loss, it=it)
logger.add('guide', 'weight_guide', weight_guide, it=it)
logger.add('losses', 'generator', gloss, it=it)
if config['training']['take_model_average']:
update_average(generator_test, generator,
beta=config['training']['model_average_beta'])
# Print stats
if it % log_every == 0:
print("weight_guide:", weight_guide) # TODO: remove
guide_loss_last = logger.get_last('losses', 'guide_loss')
g_loss_last = logger.get_last('losses', 'generator')
d_loss_last = logger.get_last('losses', 'discriminator')
d_reg_last = logger.get_last('losses', 'regularizer')
print('[epoch %0d, it %4d] g_loss = %.4f, d_loss = %.4f, reg=%.4f, guide_reg=%.4f'
% (epoch_idx, it, g_loss_last, d_loss_last, d_reg_last, guide_loss_last))
# (i) Sample if necessary
if it % config['training']['sample_every'] == 0:
print('Creating samples...')
# x = evaluator.create_samples(z_test, y_test)
x = evaluator.create_samples(z_test, clusterer.get_labels(x_test, y_test).to(device))
logger.add_imgs(x, 'all', it)
for y_inst in range(sample_nlabels):
x = evaluator.create_samples(z_test, y_inst)
logger.add_imgs(x, '%04d' % y_inst, it)
# (ii) Compute inception if necessary
if it % inception_every == 0 and it > 0:
print('PyTorch Inception score...')
inception_mean, inception_std = evaluator.compute_inception_score()
logger.add('metrics', 'pt_inception_mean', inception_mean, it=it)
logger.add('metrics', 'pt_inception_stddev', inception_std, it=it)
print(f'[epoch {epoch_idx}, it {it}] pt_inception_mean: {inception_mean}, pt_inception_stddev: {inception_std}')
# (iii) Backup if necessary
if it % backup_every == 0:
print('Saving backup...')
checkpoint_io.save('model_%08d.pt' % it, it=it)
checkpoint_io.save_clusterer(clusterer, int(it))
logger.save_stats('stats_%08d.p' % it)
if it > 0:
checkpoint_io.save('model.pt', it=it)
checkpoint_io.save('model.pt', it=it)
if __name__ == '__main__':
exit_if_job_done(out_dir)
main()
mark_job_done(out_dir)
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
