tinyflow源码笔记

std::string lua_str = "return " + lua_compute_code[inode.source->op()];
LuaRef fcompute = lua->Eval(lua_str);
op_execs_[nid] = fcompute...

struct VarState {
  /*! \brief The internal internal tensor */
  LuaRef tensor;
  /*! \brief The corresponding tblob */
  TBlob blob;

  /*! \return Whether the tensor is initialized already */
  inline bool initialized() const {
    return !tensor.is_nil();
  }
  // reset the space.
  inline void ResetSpace(TShape shape, int dev_mask = kCPU, int dtype = 0) {
    if (tensor.is_nil() ||
        shape != blob.shape ||
        dev_mask != blob.dev_mask ||
        dtype != blob.dtype) {
      TorchState* th = TorchState::ThreadLocalState();
      if (tensor.is_nil()) {
        tensor = th->NewTensorEmpty(dev_mask, dtype);
      }
      th->ResetStorage(
          tensor, th->NewStorage(shape.Size(), dev_mask, dtype), shape);
      this->blob = th->GetTBlob(tensor);
    }
  }
};

//设置参数结构
struct ConvPoolParam : public dmlc::Parameter<ConvPoolParam> {
  TShape ksize;
  TShape strides;
  std::string padding;
  std::string data_format;
  bool no_bias;
  uint32_t num_filter;

  DMLC_DECLARE_PARAMETER(ConvPoolParam) {
    DMLC_DECLARE_FIELD(ksize).set_default(TShape{1, 1, 1, 1});
    DMLC_DECLARE_FIELD(strides).set_default(TShape{1, 1, 1, 1});
    DMLC_DECLARE_FIELD(padding).set_default("SAME");
    DMLC_DECLARE_FIELD(data_format).set_default("NCHW");
    DMLC_DECLARE_FIELD(no_bias).set_default(true);
    DMLC_DECLARE_FIELD(num_filter).set_default(0);
  }
};
DMLC_REGISTER_PARAMETER(ConvPoolParam);

NNVM_REGISTER_OP(_backward)
.describe("backward operator of NN module")
.set_num_outputs([] (const NodeAttrs& attrs) {
  const NNBackwardParam& param = dmlc::get<NNBackwardParam>(attrs.parsed);
  return param.forward_readonly_inputs - param.num_no_grad_inputs;
  })
.set_num_inputs([] (const NodeAttrs& attrs) {
  const NNBackwardParam& param = dmlc::get<NNBackwardParam>(attrs.parsed);
  uint32_t n = param.num_states + 1;
  if (param.need_inputs) n += param.forward_readonly_inputs;
  if (param.need_outputs) n += 1;
  return n;
  })
.set_attr<nnvm::TIsBackward>("TIsBackward", true);

inline std::vector<NodeEntry> MakeNNBackwardNode {
  //...
  nnvm::NodePtr p = nnvm::Node::Create();
  p->attrs.op = nnvm::Op::Get("_backward");
  p->attrs.name = n->attrs.name + "_backward";
  //...
}

NNVM_REGISTER_OP_GROUP(nn_module)
.set_attr<FGradient>("FGradient", MakeNNBackwardNode)
.set_attr<bool>("TBackwardNeedInputs", true)
.set_attr<bool>("TBackwardNeedOutputs", true);


//设置infershape方法
inline bool ConvPoolShape(const NodeAttrs& attrs,
                          std::vector<TShape> *ishape,
                          std::vector<TShape> *oshape) {
  const auto& param = dmlc::get<ConvPoolParam>(attrs.parsed);
  //...
}

//应用实例
NNVM_REGISTER_OP(conv2d)
.describe("Convolution operation")
.set_num_inputs([](const NodeAttrs& attrs){
    return (dmlc::get<ConvPoolParam>(attrs.parsed).no_bias? 2 : 3);
  })
.set_attr_parser(ParamParser<ConvPoolParam>)
.include("nn_module") //上面定义的op group
.set_attr<FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
    if (dmlc::get<ConvPoolParam>(attrs.parsed).no_bias) {
      return std::vector<std::string>{"data", "weight"};
    } else {
      return std::vector<std::string>{"data", "weight", "bias"};
    }
  })
.set_attr<FInferShape>("FInferShape", ConvPoolShape)
.set_attr<bool>("TBackwardNeedOutputs", false);

NNVM_REGISTER_OP(max_pool)
.describe("Max pooling")
.set_num_inputs(1)
.set_attr_parser(ParamParser<ConvPoolParam>) //设置参数
.include("nn_module")
.set_attr<FInferShape>("FInferShape", ConvPoolShape); //设置infershape方法

NNVM_REGISTER_OP(avg_pool)
.describe("Avg pooling")
.set_num_inputs(1)
.set_attr_parser(ParamParser<ConvPoolParam>)
.include("nn_module")
.set_attr<FInferShape>("FInferShape", ConvPoolShape);

template<typename AttrType>
using FInferNodeEntryAttr = std::function<bool (const NodeAttrs& attrs,
                                                std::vector<AttrType> *in_attrs,
                                                std::vector<AttrType> *out_attrs)>;
using FInferShape = FInferNodeEntryAttr<TShape>;
using FInferType = FInferNodeEntryAttr<int>;
using TIsBackward = bool;

#define DMLC_DECLARE_PARAMETER(PType)                                   \
  static ::dmlc::parameter::ParamManager *__MANAGER__();                \
  inline void __DECLARE__(::dmlc::parameter::ParamManagerSingleton<PType> *manager) \

#define DMLC_DECLARE_FIELD(FieldName)  this->DECLARE(manager, #FieldName, FieldName)

#define DMLC_REGISTER_PARAMETER(PType)                                  \
  ::dmlc::parameter::ParamManager *PType::__MANAGER__() {               \
    static ::dmlc::parameter::ParamManagerSingleton<PType> inst(#PType); \
    return &inst.manager;                                               \
  }                                                                     \
  static DMLC_ATTRIBUTE_UNUSED ::dmlc::parameter::ParamManager&         \
  __make__ ## PType ## ParamManager__ =                                 \
      (*PType::__MANAGER__())                                           \

template<typename DType>
inline parameter::FieldEntry<DType>& DECLARE(
  parameter::ParamManagerSingleton<PType> *manager,
  const std::string &key, DType &ref) { // NOLINT(*)
    parameter::FieldEntry<DType> *e =
        new parameter::FieldEntry<DType>();
    e->Init(key, this->head(), ref);
    manager->manager.AddEntry(key, e);
    return *e;
}