如何将DNN推理阶段的所有tensor输出完整写入CSV文件

你现有代码的写入逻辑是每次迭代都创建新的DataFrame并覆盖同一份CSV文件，前序迭代的结果会被后续内容完全覆盖，因此最终仅保留最后一次迭代的输出。

推荐优先使用「先收集全量结果、最后统一写入」的方案，性能更高、避免频繁IO操作，仅在数据量极大、内存无法承载全量结果时选用逐次追加的方案。

方案1：先收集全量结果，统一写入（推荐） #

import pandas as pd
import torch
import gc

# 循环外初始化空列表，存储所有预测结果
all_predictions = []

for data in dataloader:
    # 注意：原代码中enhanced未从data取值，属于笔误，需自行补全输入读取逻辑，示例：enhanced = data[0]
    z_eval = model_back(enhanced.float().to(device))
    torch.cuda.empty_cache()
    gc.collect()

    pred = [torch.max(z.detach().cpu(), dim=1)[1] for z in z_eval]
    torch.cuda.empty_cache()
    gc.collect()
    pred_test = pred[0]
    print(pred_test)
    # 将当前批次的预测值加入总列表
    all_predictions.extend(pred_test.numpy().tolist())

# 所有迭代结束后，一次性写入CSV
pd.DataFrame(all_predictions, columns=['predictions']).to_csv('prediction.csv', index=False)

方案2：逐次追加写入（适配大数量场景） #

每次写入时指定追加模式mode='a'，仅首次写入保留表头，避免重复生成表头：

import pandas as pd
import torch
import gc

for batch_idx, data in enumerate(dataloader):
    # 注意补全enhanced的取值逻辑
    z_eval = model_back(enhanced.float().to(device))
    torch.cuda.empty_cache()
    gc.collect()

    pred = [torch.max(z.detach().cpu(), dim=1)[1] for z in z_eval]
    torch.cuda.empty_cache()
    gc.collect()
    pred_test = pred[0]
    print(pred_test)
    # 仅第一个批次写入表头，后续批次仅追加数据
    need_header = True if batch_idx == 0 else False
    pd.DataFrame(pred_test.numpy().tolist(), columns=['predictions']).to_csv(
        'prediction.csv', 
        mode='a', 
        header=need_header,
        index=False
    )

原代码循环体内的enhanced变量没有从当前迭代的data中读取输入，属于笔误漏写，需根据你自己的dataloader返回值格式补全对应取值逻辑，否则会出现输入重复、预测结果异常的问题。

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