如何降低大规模矩阵计算的内存占用？70万节点特征计算场景

针对你用A_ = torch.sigmoid(torch.matmul(x, x.t()))计算70万节点（特征维度8）两两相似度时的内存问题，半精度和分块仍无法满足需求的话，可以试试以下几个更彻底的优化方向：

• 放弃存储完整相似度矩阵，按需计算
  如果下游任务不需要完整的A_矩阵（比如只需要top-k相似节点、按批次处理下游逻辑），完全不用存储整个矩阵，而是逐批次计算并即时处理结果：

  batch_size = 1024
  for i in range(0, x.shape[0], batch_size):
      x_batch = x[i:i+batch_size]
      # 计算当前批次与所有节点的相似度
      sim_batch = torch.sigmoid(torch.matmul(x_batch, x.t()))
      # 在这里直接处理sim_batch，比如提取top-k、喂给下游任务
      # 处理完成后sim_batch会被自动释放，无需留存
  

  这种方式的内存峰值仅由batch_size决定，比如batch_size=1024时，sim_batch为1024×700000的FP16矩阵，仅占用约1.3GB，内存压力大幅降低。

• 利用低特征维度做数学简化
  你的特征维度只有8，非常小，可以拆解点积计算，避免直接生成超大中间矩阵：

  x_t = x.t()  # 转置为8×700000，方便按特征维度遍历
  sim = torch.zeros(x.shape[0], x.shape[0], dtype=torch.float16, device=x.device)
  for dim in range(8):
      dim_vec = x_t[dim:dim+1]
      # 每个特征维度的外积累加，代替一次性计算全量点积
      sim += dim_vec.t() @ dim_vec
  sim = torch.sigmoid(sim)
  

  每次仅生成一个FP16的700000×700000矩阵做累加，相比直接计算FP32中间矩阵，内存占用减半，且8次循环的额外开销几乎可以忽略。

• 用稀疏矩阵存储结果
  如果sigmoid后大部分元素值接近0（比如小于某个阈值），可以只保留大于阈值的有效元素，用稀疏矩阵格式存储：

  threshold = 0.1
  rows = []
  cols = []
  vals = []
  batch_size = 2048
  for i in range(0, x.shape[0], batch_size):
      x_batch = x[i:i+batch_size]
      sim_batch = torch.sigmoid(torch.matmul(x_batch, x.t()))
      # 过滤出有效元素
      mask = sim_batch > threshold
      row_idx, col_idx = torch.where(mask)
      row_idx += i  # 转换为全局行索引
      rows.append(row_idx)
      cols.append(col_idx)
      vals.append(sim_batch[mask])
  # 拼接为COO格式稀疏矩阵
  rows = torch.cat(rows)
  cols = torch.cat(cols)
  vals = torch.cat(vals)
  sparse_A = torch.sparse_coo_tensor(torch.stack([rows, cols]), vals, size=x.shape[:2], device=x.device)
  

  稀疏矩阵的内存占用完全取决于有效元素数量，合理设置阈值的话，内存开销能降到原有的几十分之一。

• 特征降维+哈希近似计算
  若对相似度精度要求不高，可先降维再用局部敏感哈希（LSH）分组，仅计算同组内节点的相似度：

  from torch.nn import Linear
  # 将8维特征降维到4维
  reducer = Linear(8, 4, bias=False).to(x.device)
  x_reduced = reducer(x)
  # 用简单哈希方式分组（实际可使用更专业的LSH实现）
  num_buckets = 1024
  hash_weights = torch.randn(4, num_buckets, device=x.device)
  hash_vals = torch.sum(x_reduced * hash_weights, dim=1).long() % num_buckets
  # 逐桶计算相似度
  for bucket in range(num_buckets):
      bucket_nodes = torch.where(hash_vals == bucket)[0]
      if len(bucket_nodes) == 0:
          continue
      x_bucket = x[bucket_nodes]
      sim_bucket = torch.sigmoid(torch.matmul(x_bucket, x_bucket.t()))
      # 处理当前桶内的相似度结果
  

  这种方式牺牲少量精度，但能将计算和内存开销降到接近线性水平，适合精度要求宽松的场景。

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