如何在HDFS中读取PyTorch导出的Pickle模型文件？

问题根源 #

你用hdfs dfs -get把模型文件从HDFS下载到了本地文件系统，但Spark默认会优先从HDFS路径读取文件，直接使用本地路径会导致Spark去HDFS查找对应文件，因此报出路径不存在的错误。

解决方案 #

方案1：直接从HDFS读取模型文件（推荐，无需下载到本地） #

跳过本地下载步骤，直接让Spark读取HDFS上的模型文件：

# 替换为你的HDFS模型文件完整路径，例如hdfs://xxx/xxx/model.pkl
hdfs_model_path = "hdfs://your/hdfs/path/model.pkl"
# 读取HDFS上的二进制文件
binary_data = spark.read.format("binaryFile").load(hdfs_model_path).select("content").collect()[0][0]
# 反序列化模型
import pickle
model = pickle.loads(binary_data)

方案2：读取本地已下载的模型文件 #

如果已经通过hdfs dfs -get将文件下载到本地，需要给Spark指定本地文件系统前缀file://，明确让Spark从本地读取：

import os
# 获取本地模型文件的绝对路径
local_model_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "model.pkl"))
# 添加file://前缀，告知Spark读取本地文件
spark_local_path = f"file://{local_model_path}"
# 读取本地二进制文件
binary_data = spark.read.format("binaryFile").load(spark_local_path).select("content").collect()[0][0]
# 反序列化模型
import pickle
model = pickle.loads(binary_data)

额外注意事项 #

• 确保执行任务的节点上安装的PyTorch及模型依赖库版本与训练环境一致，避免反序列化失败。
• PyTorch模型更推荐使用官方的torch.save()和torch.load()进行保存/加载，稳定性优于pickle，后续重新保存模型可参考：

  # 保存模型参数
  torch.save(model.state_dict(), "model.pt")
  # 加载模型（需先定义好模型结构）
  model = YourModelClass()
  model.load_state_dict(torch.load("model.pt"))
  

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