help='output directory for samples') parser.add_argument('--path', type=str, default='./stylegan-ffhq-1024px-new.params', help='path to checkpoint file') args = parser.parse_args() if args.gpu_id == '-1': device = mx.cpu() else: device = mx.gpu(int(args.gpu_id.strip())) generator = StyledGenerator(code_dim=512) generator.initialize() generator.collect_params().reset_ctx(device) generator.load_parameters(args.path, ctx=device) mean_style = get_mean_style(generator, device) step = int(math.log(args.size, 2)) - 2 imgs = sample(generator, step, mean_style, args.n_sample, device) if not os.path.isdir(args.out_dir): os.makedirs(args.out_dir) for i in range(args.n_sample): save_image(imgs[i], os.path.join(args.out_dir, 'sample_{}.png'.format(i)), normalize=True, img_range=(-1, 1))
'learning_rate': args.lr_default, 'beta1': 0.0, 'beta2': 0.99 }, kvstore='local') g_running = StyledGenerator(code_size) g_running.initialize(ctx=mx.gpu(0)) g_running.collect_params().reset_ctx(mx.gpu(0)) requires_grad(g_running, False) if args.ckpt_g: g_running.load_params(args.ckpt_g_running, ctx=mx.gpu(), allow_missing=True) generator.load_parameters(args.ckpt_g, ctx=context, allow_missing=True) discriminator.load_parameters(args.ckpt_d, ctx=context, allow_missing=True) accumulate(g_running, generator, 0) transform = transforms.Compose([ transforms.RandomFlipLeftRight(), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)), ]) dataset = MultiResolutionDataset(args.path, transform) if args.sched: