def _build_model(self): """Initializes the generator and discriminator networks. """ self.generator, self.discriminator = networks.build_networks( gen=self.gen_net, disc=self.disc_net, ngf=64, ndf=self.disc_size, imsize=self.imsize, nc=self.nc, k=self.noise_dim, z=self.dout_dim, bn=self.batch_norm) print('Generator', self.generator) print('Discriminator', self.discriminator) self.generator.apply(networks.weights_init) self.discriminator.apply(networks.weights_init) self.generator.cuda() self.discriminator.cuda() cudnn.benchmark = True
# Other options can change with every run parser.add_argument('--batch_size', type=int, default=64, help='Batch size [default: 64]') parser.add_argument('--fold', type=str, default='train', help='Fold [default: train]') parser.add_argument('--start_epoch', type=int, help='Epoch to start from (defaults to most recent epoch)') parser.add_argument('--count', type=int, default=1, help='Number of counterfactuals to generate') options = vars(parser.parse_args()) sys.path.append(os.path.dirname(os.path.dirname(__file__))) from dataloader import CustomDataloader import counterfactual from networks import build_networks from options import load_options # TODO: Right now, to edit cf_speed et al, you need to edit params.json start_epoch = options['start_epoch'] options = load_options(options) options['epoch'] = start_epoch dataloader = CustomDataloader(**options) # Batch size must be large enough to make a square grid visual options['batch_size'] = dataloader.num_classes + 1 networks = build_networks(dataloader.num_classes, **options) for i in range(options['count']): counterfactual.generate_counterfactual(networks, dataloader, **options)
dataloader_train = torch.utils.data.DataLoader( train_dataset, batch_size=options['batch_size'], shuffle=False, num_workers=8, pin_memory=True, drop_last=False) dataloader_off = torch.utils.data.DataLoader(dataset_off, batch_size=options['batch_size'], shuffle=False, num_workers=8, pin_memory=True, drop_last=False) networks = build_networks(**options) new_results = evaluate_with_comparison(networks, dataloader, comparison_dataloader=dataloader_off, dataloader_train=dataloader_train, **options) pprint(new_results) if not os.path.exists('evaluation'): os.mkdir('evaluation') result = {options['mode']: new_results} result_dir = os.path.join('evaluation', 'result_{}.json'.format(options['image_size'])) if os.path.exists(result_dir): old_result = json.load(open(result_dir))
help='Size of the image.') parser.add_argument( '--png', action='store_true', help='Whether to generate a png ' '(overrides num_samples and generates a 8x8 grid of images).') parser.add_argument('--o', '--out_dir', type=str, default='./', help='Output directory.') args = parser.parse_args() generator, _ = networks.build_networks(gen=args.gen, imsize=args.imsize, k=args.k) generator.load_state_dict(torch.load(args.ckpt)) print('[*] Generator loaded') if not os.path.exists(args.o): os.makedirs(args.o) generator.cuda() batch_size = 64 if args.png else 256 noise = torch.cuda.FloatTensor(batch_size, args.k, 1, 1) n_loops = 1 if args.png else args.n // batch_size + 1 samples = [] for i in tqdm(range(n_loops)): noise.normal_(0, 1) with torch.no_grad(): images = generator(noise).detach().cpu().numpy() if args.png: # Generate 8x8 if png