Function compute_feature_stats_for_generator

tools/metrics/metric_utils.py:263–321

(
    opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size: int=16,
    batch_gen=None, jit=False, temporal_detector=False, num_video_frames: int=16,
    feature_stats_cls=FeatureStats, subsample_factor: int=1, **stats_kwargs)

@torch.no_grad()
def compute_feature_stats_for_generator(
    opts, detector_url, detector_kwargs, rel_lo=0, rel_hi=1, batch_size: int=16,
    batch_gen=None, jit=False, temporal_detector=False, num_video_frames: int=16,
    feature_stats_cls=FeatureStats, subsample_factor: int=1, **stats_kwargs):

    if batch_gen is None:
        batch_gen = min(batch_size, 4)
    assert batch_size % batch_gen == 0

    # Setup generator and load labels.
    G = copy.deepcopy(opts.G).eval().requires_grad_(False).to(opts.device)
    dataset = dnnlib.util.construct_class_by_name(**opts.dataset_kwargs)

    # Image generation func.
    def run_generator(z, c, t):
        img = G(z=z, c=c, t=t, **opts.G_kwargs)
        bt, c, h, w = img.shape

        if temporal_detector:
            img = img.view(bt // num_video_frames, num_video_frames, c, h, w) # [batch_size, t, c, h, w]
            img = img.permute(0, 2, 1, 3, 4).contiguous() # [batch_size, c, t, h, w]

        # img = torch.nn.functional.interpolate(img, size=(img.shape[2], 128, 128), mode='trilinear', align_corners=False) # downsample
        # img = torch.nn.functional.interpolate(img, size=(img.shape[2], 256, 256), mode='trilinear', align_corners=False) # upsample

        img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
        return img

    # JIT.
    if jit:
        z = torch.zeros([batch_gen, G.z_dim], device=opts.device)
        c = torch.zeros([batch_gen, G.c_dim], device=opts.device)
        t = torch.zeros([batch_gen, G.cfg.sampling.num_frames_per_video], device=opts.device)
        run_generator = torch.jit.trace(run_generator, [z, c, t], check_trace=False)

    # Initialize.
    stats = feature_stats_cls(**stats_kwargs)
    assert stats.max_items is not None
    progress = opts.progress.sub(tag='generator features', num_items=stats.max_items, rel_lo=rel_lo, rel_hi=rel_hi)
    detector = get_feature_detector(url=detector_url, device=opts.device, num_gpus=opts.num_gpus, rank=opts.rank, verbose=progress.verbose)

    # Main loop.
    while not stats.is_full():
        images = []
        for _i in range(batch_size // batch_gen):
            z = torch.randn([batch_gen, G.z_dim], device=opts.device)
            cond_sample_idx = [np.random.randint(len(dataset)) for _ in range(batch_gen)]
            c = [dataset.get_label(i) for i in cond_sample_idx]
            c = torch.from_numpy(np.stack(c)).pin_memory().to(opts.device)
            t = [list(range(0, num_video_frames * subsample_factor, subsample_factor)) for _i in range(batch_gen)]
            t = torch.from_numpy(np.stack(t)).pin_memory().to(opts.device)
            images.append(run_generator(z, c, t))
        images = torch.cat(images)
        if images.shape[1] == 1:
            images = images.repeat([1, 3, *([1] * (images.ndim - 2))])
        features = detector(images, **detector_kwargs)
        stats.append_torch(features, num_gpus=opts.num_gpus, rank=opts.rank)
        progress.update(stats.num_items)