Function stack_features

src/metrics/features.py:107–142

(data_loader, eval_model, num_feats, batch_size, quantize, world_size, DDP, device, disable_tqdm)

def stack_features(data_loader, eval_model, num_feats, batch_size, quantize, world_size, DDP, device, disable_tqdm):
    eval_model.eval()
    data_iter = iter(data_loader)
    num_batches = math.ceil(float(num_feats) / float(batch_size))
    if DDP: num_batches = num_batches//world_size + 1

    real_feats, real_probs, real_labels = [], [], []
    for i in tqdm(range(0, num_batches), disable=disable_tqdm):
        start = i * batch_size
        end = start + batch_size
        try:
            images, labels = next(data_iter)
        except StopIteration:
            break

        images, labels = images.to(device), labels.to(device)

        with torch.no_grad():
            embeddings, logits = eval_model.get_outputs(images, quantize=quantize)
            probs = torch.nn.functional.softmax(logits, dim=1)
            real_feats.append(embeddings)
            real_probs.append(probs)
            real_labels.append(labels)

    real_feats = torch.cat(real_feats, dim=0)
    real_probs = torch.cat(real_probs, dim=0)
    real_labels = torch.cat(real_labels, dim=0)
    if DDP:
        real_feats = torch.cat(losses.GatherLayer.apply(real_feats), dim=0)
        real_probs = torch.cat(losses.GatherLayer.apply(real_probs), dim=0)
        real_labels = torch.cat(losses.GatherLayer.apply(real_labels), dim=0)

    real_feats = real_feats.detach().cpu().numpy().astype(np.float64)
    real_probs = real_probs.detach().cpu().numpy().astype(np.float64)
    real_labels = real_labels.detach().cpu().numpy()
    return real_feats, real_probs, real_labels