def test_case_calculate_fid_stat_CIFAR():
from template_lib.d2.data import build_dataset_mapper
from template_lib.d2template.trainer.base_trainer import build_detection_test_loader
from template_lib.v2.GAN.evaluation import build_GAN_metric
from template_lib.d2.utils.d2_utils import D2Utils
from template_lib.v2.config_cfgnode import global_cfg
from detectron2.utils import logger
logger.setup_logger('d2')
cfg = D2Utils.create_cfg()
cfg.update(global_cfg)
global_cfg.merge_from_dict(cfg)
# fmt: off
dataset_name = cfg.dataset_name
IMS_PER_BATCH = cfg.IMS_PER_BATCH
img_size = cfg.img_size
dataset_mapper_cfg = cfg.dataset_mapper_cfg
GAN_metric = cfg.GAN_metric
# fmt: on
num_workers = comm.get_world_size()
batch_size = IMS_PER_BATCH // num_workers
dataset_mapper = build_dataset_mapper(dataset_mapper_cfg, img_size=img_size)
data_loader = build_detection_test_loader(
cfg, dataset_name=dataset_name, batch_size=batch_size, mapper=dataset_mapper)
FID_IS_tf = build_GAN_metric(GAN_metric)
FID_IS_tf.calculate_fid_stat_of_dataloader(data_loader=data_loader)
comm.synchronize()
pass
def run(args):
from template_lib.d2.utils import set_ddp_seed
set_ddp_seed(outdir=f"{global_cfg.tl_outdir}/d2")
total_batch_size = global_cfg.build_dataloader.batch_size
num_workers = comm.get_world_size()
batch_size = total_batch_size // num_workers
data_loader = build_dataloader(global_cfg.build_dataloader,
kwargs_priority=True,
batch_size=batch_size,
distributed=args.distributed)
FID_IS_torch = build_GAN_metric(global_cfg.GAN_metric)
if global_cfg.tl_debug:
num_images = 50
else:
num_images = float('inf')
FID_IS_torch.calculate_fid_stat_of_dataloader(
data_loader=data_loader,
num_images=num_images,
save_fid_stat=global_cfg.save_fid_stat)
comm.synchronize()
pass
def test_case_evaluate_FID_IS():
import torch
from template_lib.v2.GAN.evaluation import build_GAN_metric
cfg_str = """
update_cfg: true
GAN_metric:
tf_fid_stat: "datasets/fid_stats_tf_cifar10.npz"
tf_inception_model_dir: "datasets/GAN_eval/tf_inception_model"
num_inception_images: 5000
"""
config = EasyDict(yaml.safe_load(cfg_str))
cfg = TFFIDISScore.update_cfg(config)
FID_IS_tf = build_GAN_metric(cfg.GAN_metric)
class SampleFunc(object):
def __init__(self, G, z):
self.G = G
self.z = z
def __call__(self, *args, **kwargs):
with torch.no_grad():
z_sample = self.z.normal_(0, 1)
# self.G.eval()
# G_z = self.G(z_sample)
G_z = self.G.normal_(0, 1)
return G_z
bs = 64
z_dim = 128
img_size = 32
z = torch.empty((bs, z_dim)).cuda()
G = torch.empty((bs, 3, img_size, img_size)).cuda()
sample_func = SampleFunc(G=G, z=z)
try:
FID_tf, IS_mean_tf, IS_std_tf = FID_IS_tf(
sample_func=sample_func, num_inception_images=5000)
print(
f'IS_mean_tf:{IS_mean_tf:.3f} +- {IS_std_tf:.3f}\n\tFID_tf: {FID_tf:.3f}'
)
except:
print("Error FID_IS_tf.")
import traceback
print(traceback.format_exc())
pass
def prepare_FID_IS(cfg):
from template_lib.v2.GAN.evaluation import build_GAN_metric
from template_lib.v2.logger import summary_dict2txtfig
from template_lib.v2.logger import global_textlogger as textlogger
logger = logging.getLogger('tl')
FID_IS = build_GAN_metric(cfg.GAN_metric)
def get_inception_metrics(sample_func, eval_iter, *args, **kwargs):
FID, IS_mean, IS_std = FID_IS(sample_func=sample_func)
logger.info(f'\n\teval_iter {eval_iter}: '
f'IS_mean_tf:{IS_mean:.3f} +- {IS_std:.3f}\n\tFID_tf: {FID:.3f}')
if not math.isnan(IS_mean):
dict_data = (dict(FID_tf=FID, IS_mean_tf=IS_mean, IS_std_tf=IS_std))
summary_dict2txtfig(dict_data=dict_data, prefix='evaltf', step=eval_iter, textlogger=textlogger)
return IS_mean, IS_std, FID
return get_inception_metrics
def test_case_evaluate_FID_IS():
import torch
from template_lib.v2.GAN.evaluation import build_GAN_metric
cfg_str = """
update_cfg: true
GAN_metric:
torch_fid_stat: "datasets/fid_stats_torch_cifar10.npz"
num_inception_images: 50000
"""
config = EasyDict(yaml.safe_load(cfg_str))
cfg = PyTorchFIDISScore.update_cfg(config)
FID_IS_torch = build_GAN_metric(cfg.GAN_metric)
class SampleFunc(object):
def __init__(self, G, z):
self.G = G
self.z = z
def __call__(self, *args, **kwargs):
with torch.no_grad():
z_sample = self.z.normal_(0, 1)
# self.G.eval()
# G_z = self.G(z_sample)
G_z = self.G.normal_(0, 1)
return G_z
bs = 64
z_dim = 128
img_size = 32
z = torch.empty((bs, z_dim)).cuda()
G = torch.empty((bs, 3, img_size, img_size)).cuda()
sample_func = SampleFunc(G=G, z=z)
FID, IS_mean, IS_std = FID_IS_torch(sample_func=sample_func)
logging.getLogger('tl').info(
f'IS_mean_tf:{IS_mean:.3f} +- {IS_std:.3f}\n\tFID_tf: {FID:.3f}')
pass
def main():
logger = logging.getLogger('tl')
args = get_args()
# CUDA setting
if not torch.cuda.is_available():
raise ValueError("Should buy GPU!")
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
device = torch.device('cuda')
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.backends.cudnn.benchmark = True
def _rescale(img):
return img * 2.0 - 1.0
def _noise_adder(img):
return torch.empty_like(img, dtype=img.dtype).uniform_(0.0,
1 / 128.0) + img
# dataset
dataset_cfg = global_cfg.get('dataset_cfg', {})
dataset_module = dataset_cfg.get('dataset_module', 'datasets.ImageFolder')
train_dataset = eval(dataset_module)(os.path.join(args.data_root),
transform=transforms.Compose([
transforms.ToTensor(),
_rescale,
_noise_adder,
]),
**dataset_cfg.get(
'dataset_kwargs', {}))
train_loader = iter(
data.DataLoader(train_dataset,
args.batch_size,
sampler=InfiniteSamplerWrapper(train_dataset),
num_workers=args.num_workers,
pin_memory=False))
if args.calc_FID:
eval_dataset = datasets.ImageFolder(
os.path.join(args.data_root, 'val'),
transforms.Compose([
transforms.ToTensor(),
_rescale,
]))
eval_loader = iter(
data.DataLoader(eval_dataset,
args.batch_size,
sampler=InfiniteSamplerWrapper(eval_dataset),
num_workers=args.num_workers,
pin_memory=True))
else:
eval_loader = None
num_classes = len(train_dataset.classes)
print(' prepared datasets...')
print(' Number of training images: {}'.format(len(train_dataset)))
# Prepare directories.
args.num_classes = num_classes
args, writer = prepare_results_dir(args)
# initialize models.
_n_cls = num_classes if args.cGAN else 0
gen_module = getattr(
global_cfg.generator, 'module',
'pytorch_sngan_projection_lib.models.generators.resnet64')
model_module = importlib.import_module(gen_module)
gen = model_module.ResNetGenerator(
args.gen_num_features,
args.gen_dim_z,
args.gen_bottom_width,
activation=F.relu,
num_classes=_n_cls,
distribution=args.gen_distribution).to(device)
if args.dis_arch_concat:
dis = SNResNetConcatDiscriminator(args.dis_num_features, _n_cls,
F.relu, args.dis_emb).to(device)
else:
dis = SNResNetProjectionDiscriminator(args.dis_num_features, _n_cls,
F.relu).to(device)
inception_model = inception.InceptionV3().to(
device) if args.calc_FID else None
opt_gen = optim.Adam(gen.parameters(), args.lr, (args.beta1, args.beta2))
opt_dis = optim.Adam(dis.parameters(), args.lr, (args.beta1, args.beta2))
# gen_criterion = getattr(L, 'gen_{}'.format(args.loss_type))
# dis_criterion = getattr(L, 'dis_{}'.format(args.loss_type))
gen_criterion = L.GenLoss(args.loss_type, args.relativistic_loss)
dis_criterion = L.DisLoss(args.loss_type, args.relativistic_loss)
print(' Initialized models...\n')
if args.args_path is not None:
print(' Load weights...\n')
prev_args, gen, opt_gen, dis, opt_dis = utils.resume_from_args(
args.args_path, args.gen_ckpt_path, args.dis_ckpt_path)
# tf FID
tf_FID = build_GAN_metric(cfg=global_cfg.GAN_metric)
class SampleFunc(object):
def __init__(self, generator, batch, latent, gen_distribution, device):
self.generator = generator
self.batch = batch
self.latent = latent
self.gen_distribution = gen_distribution
self.device = device
pass
def __call__(self, *args, **kwargs):
with torch.no_grad():
self.generator.eval()
z = utils.sample_z(self.batch, self.latent, self.device,
self.gen_distribution)
pseudo_y = utils.sample_pseudo_labels(num_classes, self.batch,
self.device)
fake_img = self.generator(z, pseudo_y)
return fake_img
sample_func = SampleFunc(gen,
batch=args.batch_size,
latent=args.gen_dim_z,
gen_distribution=args.gen_distribution,
device=device)
# Training loop
for n_iter in tqdm.tqdm(range(1, args.max_iteration + 1)):
if n_iter >= args.lr_decay_start:
decay_lr(opt_gen, args.max_iteration, args.lr_decay_start, args.lr)
decay_lr(opt_dis, args.max_iteration, args.lr_decay_start, args.lr)
# ==================== Beginning of 1 iteration. ====================
_l_g = .0
cumulative_loss_dis = .0
for i in range(args.n_dis):
if i == 0:
fake, pseudo_y, _ = sample_from_gen(args, device, num_classes,
gen)
dis_fake = dis(fake, pseudo_y)
if args.relativistic_loss:
real, y = sample_from_data(args, device, train_loader)
dis_real = dis(real, y)
else:
dis_real = None
loss_gen = gen_criterion(dis_fake, dis_real)
gen.zero_grad()
loss_gen.backward()
opt_gen.step()
_l_g += loss_gen.item()
if n_iter % 10 == 0 and writer is not None:
writer.add_scalar('gen', _l_g, n_iter)
fake, pseudo_y, _ = sample_from_gen(args, device, num_classes, gen)
real, y = sample_from_data(args, device, train_loader)
dis_fake, dis_real = dis(fake, pseudo_y), dis(real, y)
loss_dis = dis_criterion(dis_fake, dis_real)
dis.zero_grad()
loss_dis.backward()
opt_dis.step()
cumulative_loss_dis += loss_dis.item()
if n_iter % 10 == 0 and i == args.n_dis - 1 and writer is not None:
cumulative_loss_dis /= args.n_dis
writer.add_scalar('dis', cumulative_loss_dis / args.n_dis,
n_iter)
# ==================== End of 1 iteration. ====================
if n_iter % args.log_interval == 0 or n_iter == 1:
tqdm.tqdm.write(
'iteration: {:07d}/{:07d}, loss gen: {:05f}, loss dis {:05f}'.
format(n_iter, args.max_iteration, _l_g, cumulative_loss_dis))
if not args.no_image:
writer.add_image(
'fake',
torchvision.utils.make_grid(fake,
nrow=4,
normalize=True,
scale_each=True))
writer.add_image(
'real',
torchvision.utils.make_grid(real,
nrow=4,
normalize=True,
scale_each=True))
# Save previews
utils.save_images(n_iter, n_iter // args.checkpoint_interval,
args.results_root, args.train_image_root, fake,
real)
if n_iter % args.checkpoint_interval == 0:
# Save checkpoints!
utils.save_checkpoints(args, n_iter,
n_iter // args.checkpoint_interval, gen,
opt_gen, dis, opt_dis)
if (n_iter % args.eval_interval == 0
or n_iter == 1) and eval_loader is not None:
# TODO (crcrpar): implement Ineption score, FID, and Geometry score
# Once these criterion are prepared, val_loader will be used.
fid_score = evaluation.evaluate(args, n_iter, gen, device,
inception_model, eval_loader)
tqdm.tqdm.write(
'[Eval] iteration: {:07d}/{:07d}, FID: {:07f}'.format(
n_iter, args.max_iteration, fid_score))
if writer is not None:
writer.add_scalar("FID", fid_score, n_iter)
# Project embedding weights if exists.
embedding_layer = getattr(dis, 'l_y', None)
if embedding_layer is not None:
writer.add_embedding(embedding_layer.weight.data,
list(range(args.num_classes)),
global_step=n_iter)
if n_iter % global_cfg.eval_FID_every == 0 or n_iter == 1:
FID_tf, IS_mean_tf, IS_std_tf = tf_FID(sample_func=sample_func)
logger.info(
f'IS_mean_tf:{IS_mean_tf:.3f} +- {IS_std_tf:.3f}\n\tFID_tf: {FID_tf:.3f}'
)
if not math.isnan(IS_mean_tf):
summary_d = {}
summary_d['FID_tf'] = FID_tf
summary_d['IS_mean_tf'] = IS_mean_tf
summary_d['IS_std_tf'] = IS_std_tf
summary_dict2txtfig(summary_d,
prefix='train',
step=n_iter,
textlogger=global_textlogger)
gen.train()
if args.test:
shutil.rmtree(args.results_root)