from pathlib import Path
from google.protobuf import text_format
import caffe_pb2 as pb2

temp_dir = Path(".temp")
temp_dir.mkdir(exist_ok=True)

text = """
layer {
  name: "data0"
  type: "Input"
  top: "data0"
  input_param {
    shape {
      dim: 1
      dim: 3
      dim: 10
      dim: 10
    }
  }
}
layer {
  name: "output"
  type: "Concat"
  bottom: "data0"
  bottom: "data1"
  top: "output"
  concat_param {
    axis: 1
  }
}
layer {
  name: "data1"
  type: "Input"
  top: "data1"
  input_param {
    shape {
      dim: 1
      dim: 2
      dim: 10
      dim: 10
    }
  }
}

"""
predict_net = text_format.Merge(text, pb2.NetParameter())

from tvm.relax.testing import nn
from tvm.relax import op as _op
exp_tab = {} # 存储节点
dtype = "float32"
# 优先处理输入层
for pl in predict_net.layer:
    name = pl.name
    if pl.type == "Input":
        shape = pl.input_param.shape
        assert len(shape)==1 and len(pl.top)==1, "Input 类型仅仅支持单输入单输出"
        shape = list(shape[0].dim)
        exp_tab[name] = nn.Placeholder(shape, dtype, name)
for pl in predict_net.layer:
    name = pl.name
    if pl.type == "Concat":
        inputs = [exp_tab[name] for name in pl.bottom]
        exp_tab[name] = _op.concat(inputs, axis=pl.concat_param.axis)
exp_tab