| 72 | |
| 73 | |
| 74 | def prepare_evaluation(): |
| 75 | parser = ArgumentParser(add_help=True) |
| 76 | parser.add_argument("-metrics", "--eval_metrics", nargs='+', default=['fid'], |
| 77 | help="evaluation metrics to use during training, a subset list of ['fid', 'is', 'prdc'] or none") |
| 78 | parser.add_argument("--post_resizer", type=str, default="legacy", help="which resizer will you use to evaluate GANs\ |
| 79 | in ['legacy', 'clean', 'friendly']") |
| 80 | parser.add_argument('--eval_backbone', type=str, default='InceptionV3_tf',\ |
| 81 | help="[InceptionV3_tf, InceptionV3_torch, ResNet50_torch, SwAV_torch, DINO_torch, Swin-T_torch]") |
| 82 | parser.add_argument("--dset1", type=str, default=None, help="specify the directory of the folder that contains dset1 images (real).") |
| 83 | parser.add_argument("--dset1_feats", type=str, default=None, help="specify the path of *.npy that contains features of dset1 (real). \ |
| 84 | If not specified, StudioGAN will automatically extract feat1 using the whole dset1.") |
| 85 | parser.add_argument("--dset1_moments", type=str, default=None, help="specify the path of *.npy that contains moments (mu, sigma) of dset1 (real). \ |
| 86 | If not specified, StudioGAN will automatically extract moments using the whole dset1.") |
| 87 | parser.add_argument("--dset2", type=str, default=None, help="specify the directory of the folder that contains dset2 images (fake).") |
| 88 | parser.add_argument("--batch_size", default=256, type=int, help="batch_size for evaluation") |
| 89 | |
| 90 | parser.add_argument("--seed", type=int, default=-1, help="seed for generating random numbers") |
| 91 | parser.add_argument("-DDP", "--distributed_data_parallel", action="store_true") |
| 92 | parser.add_argument("--backend", type=str, default="nccl", help="cuda backend for DDP training \in ['nccl', 'gloo']") |
| 93 | parser.add_argument("-tn", "--total_nodes", default=1, type=int, help="total number of nodes for training") |
| 94 | parser.add_argument("-cn", "--current_node", default=0, type=int, help="rank of the current node") |
| 95 | parser.add_argument("--num_workers", type=int, default=8) |
| 96 | args = parser.parse_args() |
| 97 | |
| 98 | if args.dset1_feats == None and args.dset1_moments == None: |
| 99 | assert args.dset1 != None, "dset1 should be specified!" |
| 100 | if "fid" in args.eval_metrics: |
| 101 | assert args.dset1 != None or args.dset1_moments != None, "Either dset1 or dset1_moments should be given to compute FID." |
| 102 | if "prdc" in args.eval_metrics: |
| 103 | assert args.dset1 != None or args.dset1_feats != None, "Either dset1 or dset1_feats should be given to compute PRDC." |
| 104 | |
| 105 | gpus_per_node, rank = torch.cuda.device_count(), torch.cuda.current_device() |
| 106 | world_size = gpus_per_node * args.total_nodes |
| 107 | if args.seed == -1: args.seed = random.randint(1, 4096) |
| 108 | if world_size == 1: print("You have chosen a specific GPU. This will completely disable data parallelism.") |
| 109 | return args, world_size, gpus_per_node, rank |
| 110 | |
| 111 | |
| 112 | def evaluate(local_rank, args, world_size, gpus_per_node): |