Faiss通过程序加载数据到不同的显存

import faiss
import numpy as np


# 注意在启动命令里不能配置 CUDA_VISIBLE_DEVICES 
def cpu_index_to_gpu(gpu):
    ngpu = faiss.get_num_gpus()
    assert gpu < ngpu
    res = faiss.StandardGpuResources()                                  # use a single GPU
    index_flat = faiss.IndexFlatL2(32)                                  # 创建一个cpu版的Flat索引
    gpu_index_flat = faiss.index_cpu_to_gpu(res, gpu, index_flat)       # 将cpu版的索引转换成gpu版  第二个参数用于指定使用那块GPU设备
    return gpu_index_flat


# 通过  nvidia-smi 命令查看显存情况
# 加载到显存0
ins_meta = np.load('xxx0.npy')
ins_index_flat = cpu_index_to_gpu(0)
ins_index_flat.add(ins_meta)


# 加载到显存1
web_meta = np.load('xxx1.npy')
web_index_flat = cpu_index_to_gpu(1)
web_index_flat.add(web_meta)