def main():
args = parse_option()
set_seed(args.seed)
print(args)
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
# load model
assert args.netG_ckpt_step
print(f'load model from {save_path} step: {args.netG_ckpt_step}')
netG, _, netD_drs, _, _, _ = get_gan_model(
dataset_name=args.dataset,
model=args.model,
loss_type=args.loss_type,
drs=True,
)
netG.to(device)
if not args.netG_train_mode:
netG.eval()
netD_drs.eval()
netG.to(device)
netD_drs.to(device)
if args.dataset == 'celeba':
dataset = 'celeba_64'
else:
raise ValueError("Dataset should be CelebA")
evaluate_drs_with_attr(
metric='partial_recall',
attr=args.attr,
log_dir=save_path,
netG=netG,
netD_drs=netD_drs,
dataset=dataset,
num_real_samples=10000,
num_fake_samples=10000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device,
use_original_netD=args.use_original_netD,
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="color_mnist", type=str)
parser.add_argument("--root", "-r", default="./dataset/colour_mnist", type=str, help="dataset dir")
parser.add_argument("--work_dir", default="./exp_results", type=str, help="output dir")
parser.add_argument("--exp_name", default="colour_mnist", type=str, help="exp name")
parser.add_argument("--baseline_exp_name", default="colour_mnist", type=str, help="exp name")
parser.add_argument("--model", default="mnistgan", type=str, help="network model")
parser.add_argument('--gpu', default='0', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--num_pack', default=1, type=int)
parser.add_argument('--batch_size', default=64, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--num_steps', default=20000, type=int)
parser.add_argument('--logit_save_steps', default=100, type=int)
parser.add_argument('--decay', default='None', type=str)
parser.add_argument('--n_dis', default=1, type=int)
parser.add_argument('--p1_step', default=10000, type=int)
parser.add_argument('--major_ratio', default=0.99, type=float)
parser.add_argument('--num_data', default=10000, type=int)
parser.add_argument('--resample_score', type=str)
parser.add_argument("--loss_type", default="hinge", type=str, help="loss type")
parser.add_argument('--use_eval_logits', type=int)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
baseline_output_dir = f'{args.work_dir}/{args.baseline_exp_name}'
baseline_save_path = Path(baseline_output_dir)
prefix = args.exp_name.split('/')[-1]
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
netG, netD, netD_drs, optG, optD, optD_drs = get_gan_model(
dataset_name=args.dataset,
model=args.model,
drs=True,
loss_type=args.loss_type,
)
netG_ckpt_path = baseline_save_path / f'checkpoints/netG/netG_{args.p1_step}_steps.pth'
netD_ckpt_path = baseline_save_path / f'checkpoints/netD/netD_{args.p1_step}_steps.pth'
netD_drs_ckpt_path = baseline_save_path / f'checkpoints/netD/netD_{args.p1_step}_steps.pth'
logit_path = baseline_save_path / ('logits_netD_eval.pkl' if args.use_eval_logits == 1 else 'logits_netD_train.pkl')
print(f'Use logit from: {logit_path}')
logits = pickle.load(open(logit_path, "rb"))
score_start_step = args.p1_step - 5000
score_end_step = args.p1_step
score_dict = calculate_scores(logits, start_epoch=score_start_step, end_epoch=score_end_step)
sample_weights = score_dict[args.resample_score]
print(f'sample_weights mean: {sample_weights.mean()}, var: {sample_weights.var()}, max: {sample_weights.max()}, min: {sample_weights.min()}')
print_num_params(netG, netD)
ds_train = get_predefined_dataset(
dataset_name=args.dataset,
root=args.root,
weights=None,
major_ratio=args.major_ratio,
num_data=args.num_data
)
dl_train = get_dataloader(
ds_train,
batch_size=args.batch_size,
weights=sample_weights if args.resample_score is not None else None)
dl_drs = get_dataloader(ds_train, batch_size=args.batch_size, weights=None)
data_iter = iter(dl_train)
imgs, _, _, _ = next(data_iter)
plot_data(imgs, num_per_side=8, save_path=save_path, file_name=f'{prefix}_resampled_train_data_p2', vis=None)
plot_score_sort(ds_train, score_dict, save_path=save_path, phase=f'{prefix}_{score_start_step}-{score_end_step}_score', plot_metric_name=args.resample_score)
# plot_score_box(ds_train, score_dict, save_path=save_path, phase=f'{prefix}_{score_start_step}-{score_end_step}_box')
print(args, netG_ckpt_path, netD_ckpt_path, netD_drs_ckpt_path)
# Start training
trainer = LogTrainer(
output_path=save_path,
logit_save_steps=args.logit_save_steps,
netD=netD,
netG=netG,
optD=optD,
optG=optG,
netG_ckpt_file=netG_ckpt_path,
netD_ckpt_file=netD_ckpt_path,
netD_drs_ckpt_file=netD_drs_ckpt_path,
netD_drs=netD_drs,
optD_drs=optD_drs,
dataloader_drs=dl_drs,
n_dis=args.n_dis,
num_steps=args.num_steps,
save_steps=1000,
vis_steps=100,
lr_decay=args.decay,
dataloader=dl_train,
log_dir=output_dir,
print_steps=10,
device=device,
save_logits=False,
)
trainer.train()
plot_color_mnist_generator(netG, save_path=save_path, file_name=f'{prefix}-eval_p2')
netG_drs = drs.DRS(netG, netD_drs, device=device)
# for percentile in np.arange(50, 100, 5):
# netG_drs.percentile = percentile
percentile = 80
plot_color_mnist_generator(netG_drs, save_path=save_path, file_name=f'{prefix}-eval_drs_percent{percentile}_p2')
netG.restore_checkpoint(ckpt_file=netG_ckpt_path)
netG.to(device)
plot_color_mnist_generator(netG, save_path=save_path, file_name=f'{prefix}-eval_generated_p1')
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="cifar10", type=str)
parser.add_argument("--root",
"-r",
default="./dataset/cifar10",
type=str,
help="dataset dir")
parser.add_argument("--work_dir",
default="./exp_results",
type=str,
help="output dir")
parser.add_argument("--exp_name",
default="mimicry_pretrained-seed1",
type=str,
help="exp name")
parser.add_argument("--model",
default="sngan",
type=str,
help="network model")
parser.add_argument("--loss_type",
default="hinge",
type=str,
help="loss type")
parser.add_argument('--gpu',
default='0',
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--batch_size', default=128, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument("--netG_ckpt_step", type=int)
parser.add_argument("--netG_train_mode", action='store_true')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
# load model
assert args.netG_ckpt_step
print(f'load model from {save_path} step: {args.netG_ckpt_step}')
netG, _, _, _ = get_gan_model(
dataset_name=args.dataset,
model=args.model,
loss_type=args.loss_type,
)
netG.to(device)
if not args.netG_train_mode:
netG.eval()
if args.dataset == 'celeba':
dataset = 'celeba_64'
elif args.dataset == 'imagenet':
dataset = 'imagenet_128'
else:
dataset = args.dataset
if args.dataset == 'ffhq':
stats_file = './precalculated_statistics/fid_stats_ffhq_69k_run_0.npz'
# Evaluate fid
evaluate_ffhq(metric='fid',
log_dir=save_path,
data_path=args.root,
netG=netG,
dataset=dataset,
num_real_samples=50000,
num_fake_samples=50000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device,
stats_file=stats_file)
else:
if args.dataset == 'celeba':
stats_name = 'celeba_64_202k_run_0'
elif args.dataset == 'cifar10':
stats_name = 'cifar10_train'
elif args.dataset == 'imagenet':
stats_name = 'imagenet_128_50k_run_0'
stats_file = f'./precalculated_statistics/fid_stats_{stats_name}.npz'
# Evaluate fid
mmc.metrics.evaluate(metric='fid',
log_dir=save_path,
netG=netG,
dataset=dataset,
num_real_samples=50000,
num_fake_samples=50000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device,
stats_file=stats_file)
# Evaluate inception score
mmc.metrics.evaluate(metric='inception_score',
log_dir=save_path,
netG=netG,
num_samples=50000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device)
# Evaluate PR
evaluate_pr(
log_dir=save_path,
netG=netG,
dataset=dataset,
num_real_samples=10000,
num_fake_samples=10000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device,
)
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="cifar10", type=str)
parser.add_argument("--root", "-r", default="./dataset/cifar10", type=str, help="dataset dir")
parser.add_argument("--work_dir", default="./exp_results", type=str, help="output dir")
parser.add_argument("--exp_name", default="mimicry_pretrained-seed1", type=str, help="exp name")
parser.add_argument('--gpu', default='0', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--batch_size', default=128, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--epochs', default=50, type=int)
parser.add_argument("--netG_step", type=int)
parser.add_argument("--netG_train_mode", action='store_true')
parser.add_argument("--cae_ckpt_path", type=str)
parser.add_argument("--model", type=str)
parser.add_argument("--loss_type", default='ns', type=str)
parser.add_argument("--generated_dataset_path", type=str)
parser.add_argument('--major_ratio', default=0.99, type=float)
parser.add_argument('--num_data', default=10000, type=int)
parser.add_argument('--num_pack', default=1, type=int)
parser.add_argument('--topk', action='store_true')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
if args.dataset == 'mnist_c':
ds_test = get_predefined_dataset(
dataset_name=args.dataset,
root=args.root,
)
else:
ds_test = get_predefined_dataset(
dataset_name=args.dataset,
root=args.root,
major_ratio=args.major_ratio,
num_data=args.num_data
)
dl_test = get_dataloader(dataset=ds_test, batch_size=args.batch_size)
# load model
assert args.netG_step
print(f'load model from: {args.netG_step}')
netG, _, netD_drs, _, _, _ = get_gan_model(
args.dataset,
model=args.model,
drs=True,
loss_type=args.loss_type,
topk=args.topk,
num_pack=args.num_pack,
inclusive=True,
num_data=args.num_data,
dataloader=dl_test,)
netG.to(device)
netD_drs.to(device)
netG.get_setting(train=False)
step = netG.restore_checkpoint(ckpt_file=save_path / f'checkpoints/netG/netG_{args.netG_step}_steps.pth')
netD_drs_ckpt_path = save_path / f'checkpoints/netD_drs/netD_drs_{args.netG_step}_steps.pth'
if os.path.exists(netD_drs_ckpt_path):
use_drs = True
netD_drs.restore_checkpoint(ckpt_file=netD_drs_ckpt_path)
netD_drs.to(device)
drs = DRS(netG=netG, netD=netD_drs, device=device)
else:
use_drs = False
drs = netG
print(f'use drs: {use_drs}')
model = get_ae_model(dataset_name=args.dataset).to(device)
if args.cae_ckpt_path:
model.load_state_dict(torch.load(args.cae_ckpt_path))
else:
if args.generated_dataset_path:
print(f'skip data generation, use: {args.generated_dataset_path}')
generated_dataset_path = args.generated_dataset_path
else:
# generate dataset
generated_dataset_path = save_path / f'netG_{step}_steps_seed{args.seed}_generated_dataset.pkl'
generate_dataset(drs, generated_dataset_path, eval_mode=not args.netG_train_mode, device=device)
print(f'data generated in: {generated_dataset_path}')
ds_train = get_generated_dataset(dataset_name=args.dataset, root=generated_dataset_path)
dl_train = get_dataloader(dataset=ds_train, batch_size=args.batch_size)
cae_ckpt_path = save_path / 'cae_checkpoints' / f'{step}_steps_seed{args.seed}'
cae_ckpt_path.mkdir(parents=True, exist_ok=True)
model = train_cae(model, dl_train=dl_train, dl_test=dl_test, save_path=cae_ckpt_path, epochs=args.epochs)
final_loss = test_cae(dl_test, model)
final_score = final_loss
pickle.dump(final_score, open(save_path / f'netG_{step}_steps_seed{args.seed}_epoch{args.epochs}_ae_score.pkl', 'wb'))
show_sorted_score_samples(
dataset=ds_test,
score=final_score,
save_path=save_path,
score_name='ae_score',
plot_name=f'netG_{step}_steps_seed{args.seed}_epoch{args.epochs}_ae_score')
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="color_mnist", type=str)
parser.add_argument("--root",
"-r",
default="./dataset/colour_mnist",
type=str,
help="dataset dir")
parser.add_argument("--work_dir",
default="./exp_results",
type=str,
help="output dir")
parser.add_argument("--exp_name",
default="colour_mnist",
type=str,
help="exp name")
parser.add_argument("--loss_type",
default="ns",
type=str,
help="loss type")
parser.add_argument("--model",
default="mnist_dcgan",
type=str,
help="network model")
parser.add_argument('--gpu',
default='0',
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--num_pack', default=1, type=int)
parser.add_argument('--batch_size', default=64, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--use_clipping', action='store_true')
parser.add_argument('--num_steps', default=20000, type=int)
parser.add_argument('--logit_save_steps', default=100, type=int)
parser.add_argument('--decay', default='None', type=str)
parser.add_argument('--n_dis', default=1, type=int)
parser.add_argument('--major_ratio', default=0.99, type=float)
parser.add_argument('--num_data', default=10000, type=int)
parser.add_argument('--topk', default=0, type=int)
parser.add_argument('--resample_score', type=str)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
ds_train = get_predefined_dataset(dataset_name=args.dataset,
root=args.root,
weights=None,
major_ratio=args.major_ratio,
num_data=args.num_data)
dl_train = get_dataloader(ds_train,
batch_size=args.batch_size,
weights=None)
netG, netD, optG, optD = get_gan_model(
dataset_name=args.dataset,
model=args.model,
num_pack=args.num_pack,
loss_type=args.loss_type,
topk=args.topk == 1,
inclusive=True,
num_data=args.num_data,
dataloader=dl_train,
)
print_num_params(netG, netD)
print(args)
# Start training
trainer = LogTrainer(
output_path=save_path,
logit_save_steps=args.logit_save_steps,
netD=netD,
netG=netG,
optD=optD,
optG=optG,
n_dis=args.n_dis,
num_steps=args.num_steps,
save_steps=1000,
vis_steps=100,
lr_decay=args.decay,
dataloader=dl_train,
log_dir=output_dir,
print_steps=10,
device=device,
topk=args.topk,
save_logits=args.num_pack == 1,
save_eval_logits=False,
)
trainer.train()
plot_color_mnist_generator(netG, save_path=save_path, file_name='eval_p1')
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="color_mnist", type=str)
parser.add_argument("--root", "-r", default="./dataset/colour_mnist", type=str, help="dataset dir")
parser.add_argument("--work_dir", default="./exp_results", type=str, help="output dir")
parser.add_argument("--exp_name", default="colour_mnist", type=str, help="exp name")
parser.add_argument("--baseline_exp_name", default="colour_mnist", type=str, help="exp name")
parser.add_argument("--model", default="mnistgan", type=str, help="network model")
parser.add_argument('--gpu', default='0', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--num_pack', default=1, type=int)
parser.add_argument('--batch_size', default=64, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--use_clipping', action='store_true')
parser.add_argument('--num_steps', default=20000, type=int)
parser.add_argument('--logit_save_steps', default=100, type=int)
parser.add_argument('--decay', default='None', type=str)
parser.add_argument('--n_dis', default=1, type=int)
parser.add_argument('--p1_step', default=10000, type=int)
parser.add_argument('--major_ratio', default=0.99, type=float)
parser.add_argument('--num_data', default=10000, type=int)
parser.add_argument('--resample_score', type=str)
parser.add_argument("--loss_type", default="hinge", type=str, help="loss type")
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
baseline_output_dir = f'{args.work_dir}/{args.baseline_exp_name}'
baseline_save_path = Path(baseline_output_dir)
prefix = args.exp_name.split('/')[-1]
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
netG, netD, optG, optD = get_gan_model(
dataset_name=args.dataset,
model=args.model,
loss_type=args.loss_type,
gold=True
)
netG_ckpt_path = baseline_save_path / f'checkpoints/netG/netG_{args.p1_step}_steps.pth'
netD_ckpt_path = baseline_save_path / f'checkpoints/netD/netD_{args.p1_step}_steps.pth'
print_num_params(netG, netD)
ds_train = get_predefined_dataset(
dataset_name=args.dataset,
root=args.root,
weights=None,
major_ratio=args.major_ratio,
num_data=args.num_data
)
dl_train = get_dataloader(
ds_train,
batch_size=args.batch_size,
weights=None)
data_iter = iter(dl_train)
imgs, _, _, _ = next(data_iter)
plot_data(imgs, num_per_side=8, save_path=save_path, file_name=f'{prefix}_gold_train_data_p2', vis=None)
print(args, netG_ckpt_path, netD_ckpt_path)
# Start training
trainer = LogTrainer(
output_path=save_path,
logit_save_steps=args.logit_save_steps,
netD=netD,
netG=netG,
optD=optD,
optG=optG,
netG_ckpt_file=netG_ckpt_path,
netD_ckpt_file=netD_ckpt_path,
n_dis=args.n_dis,
num_steps=args.num_steps,
save_steps=1000,
vis_steps=100,
lr_decay=args.decay,
dataloader=dl_train,
log_dir=output_dir,
print_steps=10,
device=device,
save_logits=False,
gold=True,
gold_step=args.p1_step
)
trainer.train()
plot_color_mnist_generator(netG, save_path=save_path, file_name=f'{prefix}-eval_p2')
netG.restore_checkpoint(ckpt_file=netG_ckpt_path)
netG.to(device)
plot_color_mnist_generator(netG, save_path=save_path, file_name=f'{prefix}-eval_generated_p1')
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="cifar10", type=str)
parser.add_argument("--root",
"-r",
default="./dataset/cifar10",
type=str,
help="dataset dir")
parser.add_argument("--work_dir",
default="./exp_results",
type=str,
help="output dir")
parser.add_argument("--exp_name", type=str, help="exp name")
parser.add_argument("--baseline_exp_name", type=str, help="exp name")
parser.add_argument('--p1_step', default=40000, type=int)
parser.add_argument("--model",
default="sngan",
type=str,
help="network model")
parser.add_argument("--loss_type",
default="hinge",
type=str,
help="loss type")
parser.add_argument('--gpu',
default='0',
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--num_steps', default=80000, type=int)
parser.add_argument('--batch_size', default=64, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--decay', default='linear', type=str)
parser.add_argument('--n_dis', default=5, type=int)
parser.add_argument('--resample_score', type=str)
parser.add_argument('--gold', action='store_true')
parser.add_argument('--topk', action='store_true')
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
baseline_output_dir = f'{args.work_dir}/{args.baseline_exp_name}'
baseline_save_path = Path(baseline_output_dir)
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
prefix = args.exp_name.split('/')[-1]
if args.dataset == 'celeba':
window = 5000
elif args.dataset == 'cifar10':
window = 5000
else:
window = 5000
if not args.gold:
logit_path = baseline_save_path / 'logits_netD_eval.pkl'
print(f'Use logit from: {logit_path}')
logits = pickle.load(open(logit_path, "rb"))
score_start_step = (args.p1_step - window)
score_end_step = args.p1_step
score_dict = calculate_scores(logits,
start_epoch=score_start_step,
end_epoch=score_end_step)
sample_weights = score_dict[args.resample_score]
print(
f'sample_weights mean: {sample_weights.mean()}, var: {sample_weights.var()}, max: {sample_weights.max()}, min: {sample_weights.min()}'
)
else:
sample_weights = None
netG_ckpt_path = baseline_save_path / f'checkpoints/netG/netG_{args.p1_step}_steps.pth'
netD_ckpt_path = baseline_save_path / f'checkpoints/netD/netD_{args.p1_step}_steps.pth'
netD_drs_ckpt_path = baseline_save_path / f'checkpoints/netD/netD_{args.p1_step}_steps.pth'
netG, netD, netD_drs, optG, optD, optD_drs = get_gan_model(
dataset_name=args.dataset,
model=args.model,
loss_type=args.loss_type,
drs=True,
topk=args.topk,
gold=args.gold,
)
print(f'model: {args.model} - netD_drs_ckpt_path: {netD_drs_ckpt_path}')
print_num_params(netG, netD)
ds_train = get_predefined_dataset(dataset_name=args.dataset,
root=args.root,
weights=None)
dl_train = get_dataloader(ds_train,
batch_size=args.batch_size,
weights=sample_weights)
ds_drs = get_predefined_dataset(dataset_name=args.dataset,
root=args.root,
weights=None)
dl_drs = get_dataloader(ds_drs, batch_size=args.batch_size, weights=None)
if not args.gold:
show_sorted_score_samples(ds_train,
score=sample_weights,
save_path=save_path,
score_name=args.resample_score,
plot_name=prefix)
print(args)
# Start training
trainer = LogTrainer(
output_path=save_path,
netD=netD,
netG=netG,
optD=optD,
optG=optG,
netG_ckpt_file=str(netG_ckpt_path),
netD_ckpt_file=str(netD_ckpt_path),
netD_drs_ckpt_file=str(netD_drs_ckpt_path),
netD_drs=netD_drs,
optD_drs=optD_drs,
dataloader_drs=dl_drs,
n_dis=args.n_dis,
num_steps=args.num_steps,
save_steps=1000,
lr_decay=args.decay,
dataloader=dl_train,
log_dir=output_dir,
print_steps=10,
device=device,
topk=args.topk,
gold=args.gold,
gold_step=args.p1_step,
save_logits=False,
)
trainer.train()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="cifar10", type=str)
parser.add_argument("--work_dir", default="./exp_results", type=str, help="output dir")
parser.add_argument("--exp_name", default="mimicry_pretrained-seed1", type=str, help="exp name")
parser.add_argument("--baseline_exp_name", type=str, help="exp name")
parser.add_argument('--p1_step', default=40000, type=int)
parser.add_argument("--model", default="sngan", type=str, help="network model")
parser.add_argument("--loss_type", default="hinge", type=str, help="loss type")
parser.add_argument('--gpu', default='0', type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--batch_size', default=128, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument("--netG_ckpt_step", type=int)
parser.add_argument("--netG_train_mode", action='store_true')
parser.add_argument('--resample_score', type=str)
parser.add_argument('--gold', action='store_true')
parser.add_argument('--topk', action='store_true')
parser.add_argument("--index_num", default=100, type=int, help="number of index to use for FID score")
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
baseline_output_dir = f'{args.work_dir}/{args.baseline_exp_name}'
baseline_save_path = Path(baseline_output_dir)
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
# load model
assert args.netG_ckpt_step
print(f'load model from {save_path} step: {args.netG_ckpt_step}')
netG, _, _, _ = get_gan_model(
dataset_name=args.dataset,
model=args.model,
loss_type=args.loss_type,
topk=args.topk,
gold=args.gold,
)
netG.to(device)
if not args.netG_train_mode:
netG.eval()
if args.dataset == 'celeba':
dataset = 'celeba_64'
window = 5000
else:
dataset = args.dataset
window = 5000
logit_path = baseline_save_path / 'logits_netD_eval.pkl'
print(f'Use logit from: {logit_path}')
logits = pickle.load(open(logit_path, "rb"))
score_start_step = (args.p1_step - window)
score_end_step = args.p1_step
score_dict = calculate_scores(logits, start_epoch=score_start_step, end_epoch=score_end_step)
sample_weights = score_dict[args.resample_score]
print(
f'sample_weights mean: {sample_weights.mean()}, var: {sample_weights.var()}, max: {sample_weights.max()}, min: {sample_weights.min()}')
print(args)
sort_index = np.argsort(sample_weights)
high_index = sort_index[-args.index_num:]
low_index = sort_index[:args.index_num]
# Evaluate fid with index of high weight
evaluate_with_index(
metric='fid',
index=high_index,
log_dir=save_path,
netG=netG,
dataset=dataset,
num_fake_samples=50000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device,
stats_file=None,
name=f'high_{args.resample_score}', )
# Evaluate fid with index of low weight
evaluate_with_index(
metric='fid',
index=low_index,
log_dir=save_path,
netG=netG,
dataset=dataset,
num_fake_samples=50000,
evaluate_step=args.netG_ckpt_step,
num_runs=1,
device=device,
stats_file=None,
name=f'low_{args.resample_score}', )
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--dataset", "-d", default="cifar10", type=str)
parser.add_argument("--root",
"-r",
default="./dataset/cifar10",
type=str,
help="dataset dir")
parser.add_argument("--work_dir",
default="./exp_results",
type=str,
help="output dir")
parser.add_argument("--exp_name",
default="cifar10",
type=str,
help="exp name")
parser.add_argument("--model",
default="sngan",
type=str,
help="network model")
parser.add_argument("--loss_type",
default="hinge",
type=str,
help="loss type")
parser.add_argument('--gpu',
default='0',
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--num_pack', default=1, type=int)
parser.add_argument('--batch_size', default=64, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument('--download_dataset', action='store_true')
parser.add_argument('--topk', action='store_true')
parser.add_argument('--num_steps', default=100000, type=int)
parser.add_argument('--logit_save_steps', default=100, type=int)
parser.add_argument('--decay', default='linear', type=str)
parser.add_argument('--n_dis', default=5, type=int)
parser.add_argument('--imb_factor', default=0.1, type=float)
parser.add_argument('--celeba_class_attr', default='glass', type=str)
parser.add_argument('--ckpt_step', type=int)
parser.add_argument('--no_save_logits', action='store_true')
parser.add_argument('--save_logit_after', default=30000, type=int)
parser.add_argument('--stop_save_logit_after', default=60000, type=int)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
output_dir = f'{args.work_dir}/{args.exp_name}'
save_path = Path(output_dir)
save_path.mkdir(parents=True, exist_ok=True)
set_seed(args.seed)
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
netG, netD, optG, optD = get_gan_model(
dataset_name=args.dataset,
model=args.model,
loss_type=args.loss_type,
topk=args.topk,
)
print_num_params(netG, netD)
ds_train = get_predefined_dataset(
dataset_name=args.dataset,
root=args.root,
)
dl_train = get_dataloader(ds_train, batch_size=args.batch_size)
if args.dataset == 'celeba':
args.num_steps = 75000
args.logit_save_steps = 100
args.save_logit_after = 55000
args.stop_save_logit_after = 60000
if args.dataset == 'cifar10':
args.num_steps = 50000
args.logit_save_steps = 100
args.save_logit_after = 35000
args.stop_save_logit_after = 40000
print(args)
if args.ckpt_step:
netG_ckpt_file = save_path / f'checkpoints/netG/netG_{args.ckpt_step}_steps.pth'
netD_ckpt_file = save_path / f'checkpoints/netD/netD_{args.ckpt_step}_steps.pth'
else:
netG_ckpt_file = None
netD_ckpt_file = None
# Start training
trainer = LogTrainer(
output_path=save_path,
logit_save_steps=args.logit_save_steps,
netG_ckpt_file=netG_ckpt_file,
netD_ckpt_file=netD_ckpt_file,
netD=netD,
netG=netG,
optD=optD,
optG=optG,
n_dis=args.n_dis,
num_steps=args.num_steps,
save_steps=1000,
lr_decay=args.decay,
dataloader=dl_train,
log_dir=output_dir,
print_steps=10,
device=device,
topk=args.topk,
save_logits=not args.no_save_logits,
save_logit_after=args.save_logit_after,
stop_save_logit_after=args.stop_save_logit_after,
)
trainer.train()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--work_dir",
default="./exp_results",
type=str,
help="output dir")
parser.add_argument("--exp_name",
default="mimicry_pretrained-seed1",
type=str,
help="exp name")
parser.add_argument("--model",
default="sngan",
type=str,
help="network model")
parser.add_argument("--loss_type",
default="hinge",
type=str,
help="loss type")
parser.add_argument("--classifier",
default="vgg16",
type=str,
help="calssifier network model")
parser.add_argument('--gpu',
default='0',
type=str,
help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--batch_size', default=100, type=int)
parser.add_argument('--seed', default=1, type=int)
parser.add_argument("--netG_ckpt_step", type=int)
parser.add_argument("--netG_train_mode", action='store_true')
parser.add_argument("--use_original_netD", action='store_true')
parser.add_argument('--attr', default='Bald', type=str)
parser.add_argument('--drs', action='store_true')
parser.add_argument('--num_samples', default=50000, type=int)
args = parser.parse_args()
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu
set_seed(args.seed)
print(args)
save_path = f'{args.work_dir}/{args.exp_name}'
if torch.cuda.is_available():
device = "cuda"
cudnn.benchmark = True
else:
device = "cpu"
# load model
assert args.netG_ckpt_step
print(f'load model from {save_path} step: {args.netG_ckpt_step}')
if args.drs:
netG, _, netD_drs, _, _, _ = get_gan_model(dataset_name='celeba',
model=args.model,
loss_type=args.loss_type,
drs=True)
else:
netG, _, _, _ = get_gan_model(
dataset_name='celeba',
model=args.model,
loss_type=args.loss_type,
)
netG.to(device)
if not args.netG_train_mode:
netG.eval()
netG.to(device)
if args.drs:
netD_drs.eval()
netD_drs.to(device)
gan_ckpt = f'{args.work_dir}/{args.exp_name}/checkpoints/netG/netG_{args.netG_ckpt_step}_steps.pth'
if args.use_original_netD:
netD_drs_ckpt = f'{args.work_dir}/{args.exp_name}/checkpoints/netD/netD_{args.netG_ckpt_step}_steps.pth'
else:
netD_drs_ckpt = f'{args.work_dir}/{args.exp_name}/checkpoints/netD_drs/netD_drs_{args.netG_ckpt_step}_steps.pth'
print(gan_ckpt)
netG.restore_checkpoint(ckpt_file=gan_ckpt)
if args.drs:
netD_drs.restore_checkpoint(ckpt_file=netD_drs_ckpt)
netG = DRS(netG=netG, netD=netD_drs, device=device)
# load classifier
print('Load classifier')
if args.classifier == 'vgg16':
model = models.vgg16(pretrained=True)
elif args.classifier == 'resnet18':
model = models.resnet18(pretrained=True)
elif args.classifier == 'inception':
model = models.inception_v3(pretrained=True)
else:
raise ValueError('model should be vgg16 or resnet18 or inception')
# change the number of classes
in_features = model.classifier[6].in_features
model.classifier[6] = nn.Linear(in_features, 2, bias=True)
classifier_path = './convnet_celeba'
model.load_state_dict(
torch.load(os.path.join(classifier_path, f'{args.attr}.pth')))
model.to(device)
batch_size = min(args.batch_size, args.num_samples)
num_batches = args.num_samples // batch_size
attr_num = 0
not_attr_num = 0
for i in range(num_batches):
with torch.no_grad():
img = netG.generate_images(batch_size, device=device)
labels = model(img)
answers = torch.argmax(labels, dim=1)
attr = torch.count_nonzero(answers).item()
not_attr = batch_size - attr
attr_num += attr
not_attr_num += not_attr
print(f'attr: {attr_num}')
print(f'not attr: {not_attr_num}')
output_dir = os.path.join(save_path, 'evaluate',
f'step-{args.netG_ckpt_step}')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_file = os.path.join(output_dir, f'count_attribute.csv')
if os.path.exists(output_file):
with open(output_file, 'a', newline='') as f:
wr = csv.writer(f)
wr.writerow([args.attr, attr_num, not_attr_num])
else:
with open(output_file, 'w', newline='') as f:
wr = csv.writer(f)
wr.writerow(['', 'attr', 'not attr'])
wr.writerow([args.attr, attr_num, not_attr_num])