Go语言基于graft库从内存加载TensorFlow模型的可行方案咨询

你目前用Go语言结合graft库调用TensorFlow，通过文件路径加载模型，想要实现从内存直接加载的需求，TensorFlow官方Go绑定本身没有提供对应的LoadModelFromMemory接口，但有两种无需修改TensorFlow源码的现成方案：

方案一：临时文件中转（最简便可行） #

把内存中的模型字节写入临时文件，再用tf.LoadSavedModel（graft库底层依赖的TensorFlow方法）加载，加载完成后删除临时文件即可。这种方式不需要改动任何底层代码，兼容性拉满。

示例代码：

import (
    "bytes"
    "io"
    "os"
    "path/filepath"
    "archive/zip"
    tf "github.com/tensorflow/tensorflow/tensorflow/go"
)

// LoadModelFromMemory 从内存字节加载SavedModel
func LoadModelFromMemory(modelBytes []byte, tags []string) (*tf.SavedModel, error) {
    // 创建临时目录
    tempDir, err := os.MkdirTemp("", "tf-mem-model-*")
    if err != nil {
        return nil, err
    }
    // 延迟清理临时文件，即使加载失败也会执行
    defer func() { _ = os.RemoveAll(tempDir) }()

    // 假设modelBytes是SavedModel目录的zip压缩包，先解压到临时目录
    zipReader, err := zip.NewReader(bytes.NewReader(modelBytes), int64(len(modelBytes)))
    if err != nil {
        return nil, err
    }
    for _, file := range zipReader.File {
        filePath := filepath.Join(tempDir, file.Name)
        // 创建文件所在目录
        if err := os.MkdirAll(filepath.Dir(filePath), 0755); err != nil {
            return nil, err
        }
        // 写入文件
        dstFile, err := os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, file.Mode())
        if err != nil {
            return nil, err
        }
        fileIn, err := file.Open()
        if err != nil {
            _ = dstFile.Close()
            return nil, err
        }
        _, err = io.Copy(dstFile, fileIn)
        _ = fileIn.Close()
        _ = dstFile.Close()
        if err != nil {
            return nil, err
        }
    }

    // 如果是单个.pb冻结模型，替换上面的解压逻辑为：
    // modelFile := filepath.Join(tempDir, "frozen_model.pb")
    // if err := os.WriteFile(modelFile, modelBytes, 0644); err != nil {
    //     return nil, err
    // }
    // 后续加载冻结模型的逻辑需对应调整

    // 加载SavedModel
    model, err := tf.LoadSavedModel(tempDir, tags, nil)
    if err != nil {
        return nil, err
    }
    return model, nil
}

如果使用graft库，只需将加载后的tf.SavedModel传入graft的相关初始化方法即可，逻辑完全兼容。

方案二：封装TensorFlow C API内存加载接口 #

TensorFlow的C API提供了TF_LoadSessionFromMemory这类直接从内存加载模型的接口，但Go绑定没有暴露。你可以通过cgo直接调用这些C接口，封装成Go可用的函数。这种方式避免了磁盘IO，但需要编写cgo代码，复杂度比临时文件方案高，适合对性能要求极高的场景。

内容的提问来源于stack exchange，提问作者José Ángel Martínez