如何在不拆分的前提下对超大型Torch张量执行运算

我正尝试为k近邻（k-Nearest-Neighbours）算法计算两个大型张量中每对样本的两两距离：给定形状为(b1,c,h,w)的test张量与形状为(b2,c,h,w)的train张量，需要对所有i、j计算|| test[i]-train[j] ||（test[i]与train[j]均为批次内样本，形状为(c,h,w)）。

train与test张量体量极大，无法完整载入RAM（内存）。

我未一次性构建全量张量，而是在构建过程中将数据张量拆分后分别存储，最终得到分片文件集合{Test\test_1,...,Test\test_n}与{Train\train_1,...,Train\train_m}；之后通过嵌套for循环逐次加载每个Test\test_i与Train\train_j分片，计算对应距离结果后存储。

对应半伪代码实现如下：

test_files = [f'Test\test_{i}' for i in range(n)]
train_files = [f'Train\train_{j}' for j in range(m)]
dist = lambda t1,t2: torch.cdist(t1.flatten(1), t2.flatten(1)) 
all_distances = []
for test_i in test_files:
    test_i = torch.load(test_i) # 形状为(c,h,w)
    dist_of_i_from_all_j = torch.Tensor([])
    for train_j in train_files:
        train_j = torch.load(train_j) # 形状为(c,h,w)
        dist_of_i_from_all_j = torch.cat((dist_of_i_from_all_j, dist(test_i,train_j))
    all_distances.append(dist_of_i_from_all_j)
# 最终从all_distances中取k个最小值即可

我了解到FAISS提供的方案或许可以加速该流程，但不清楚具体落地方法，任何可行的实现方案都可作为参考。

内容的提问来源于stack exchange，提问作者Hadar