INNER CODE UNIT · Python

generate_embeddings

openfoodfacts/openfoodfacts-ai · logo-ann/benchmarks/embedding_models_benchmark/main.py:75

def generate_embeddings(
    model: torch.nn.Module, data_loader: DataLoader, device: torch.device
) -> Tuple[torch.Tensor, torch.Tensor, float]:
    model.eval()
    embedding_all = []
    labels_all = []
    elapsed = 0.0

    with torch.inference_mode():
        for inputs, labels in tqdm.tqdm(data_loader):
            start_time = time.monotonic()
            output = model(inputs.to(device))
            elapsed += time.monotonic() - start_time
            embedding_all.append(output)
            labels_all.append(labels)

        return torch.cat(embedding_all), torch.cat(labels_all), elapsed

View source record →