当然可以！不过你当前的代码用了with tf.Session() as sess:的写法，这种情况下with代码块结束后会话会自动关闭，直接返回的话拿到的是已经关闭的会话，根本没法复用。所以得调整会话的创建方式：

修改后的代码示例 #

def train_model(X_data, y_data):
    # 假设你已经在函数内（或外部）定义了X、y占位符、optimizer优化器和init初始化操作
    # 手动创建会话，不使用with语句
    sess = tf.Session()
    # 初始化变量
    sess.run(init)
    # 执行训练循环
    for epoch in range(200):
        for (xh, yh) in zip(X_data, y_data):
            sess.run(optimizer, feed_dict={X: xh, y: yh})
    # 返回训练好的会话
    return sess

关键注意事项 #

• 手动关闭会话：因为没有用with自动管理，你在使用完返回的会话后，必须手动调用sess.close()来释放资源，避免内存泄漏。比如：

  trained_sess = train_model(X_train, y_train)
  # 用会话做预测或其他操作
  predictions = trained_sess.run(y_pred, feed_dict={X: test_data})
  # 用完记得关闭
  trained_sess.close()
  

• 保持计算图一致：确保你后续使用会话时，所有操作（比如预测用的y_pred）都在同一个TensorFlow计算图中。如果你的占位符、模型操作是在train_model函数内定义的，那后续使用也要基于这个图；如果是在函数外定义的全局图，就没问题。
• TF版本提示：如果你用的是TensorFlow 2.x，原生已经废弃了Session这种显式会话的写法，转而使用Eager Execution或tf.function。如果是新项目，建议迁移到TF2.x的风格，但如果是维护旧代码，上面的方法完全适用。

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