Error Caffe2ModelLoader::loadWeight()

in lib/Importer/Caffe2ModelLoader.cpp [2455:2923]


Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
  ArgumentDictionaryTy dict = loadArgumentMap(op);
  const std::string &typeName = op.type();
  const std::string &opName = loadOperatorName(op);
  // Load tensors with values:
  if (typeName == "GivenTensorFill" || typeName == "GivenTensorFp16Fill" ||
      typeName == "GivenTensorIntFill" || typeName == "GivenTensorInt64Fill") {
    /*
     * op {
     *   output: "conv1_w"
     *   name: ""
     *   type: "GivenTensorFill"
     *   arg {
     *     name: "shape"
     *     ints: 96
     *     ints: 3
     *     ints: 11
     *     ints: 11
     *   }
     *   arg {
     *     name: "values"
     *     floats: -0.028315347
     *     ...
     *   }
     * }
     */

    // Note: Explicitly allow for an empty dim here, representing a scalar value
    // will be loaded below.
    std::vector<dim_t> dim;
    ASSIGN_VALUE_OR_RETURN_ERR(
        dim, getShape<dim_t>(dict["shape"], /* allowEmptyShape */ true));
    auto const &values = dict["values"];
    RETURN_ERR_IF_NOT(
        op.output_size() == 1,
        opErrMsg(
            op, strFormat(
                    "GivenTensorFill must have exactly 1 output, but found %d ",
                    op.output_size())));
    Tensor T;
    if (typeName == "GivenTensorFill") {
      RETURN_IF_ERR(
          fillTensor<float>(T, ElemKind::FloatTy, dim, values->floats()));
    } else if (typeName == "GivenTensorFp16Fill") {
      RETURN_IF_ERR(
          fillTensor<float16_t>(T, ElemKind::Float16Ty, dim, values->floats()));
    } else if (typeName == "GivenTensorIntFill") {
      RETURN_IF_ERR(
          fillTensor<int32_t>(T, ElemKind::Int32ITy, dim, values->ints()));
    } else if (typeName == "GivenTensorInt64Fill") {
      RETURN_IF_ERR(
          fillTensor<int64_t>(T, ElemKind::Int64ITy, dim, values->ints()));
    } else {
      return MAKE_ERR(
          strFormat("Unhandled tensor fill type: %s", typeName.c_str()));
    }
    RETURN_IF_ERR(createAndRegisterConstant(op.output().Get(0), std::move(T)));
    return Error::success();
  }

  if (typeName == "GivenTensorByteStringToUInt8Fill") {
    /*
      output: "data"
      type: "GivenTensorByteStringToUInt8Fill"
      arg {
      name: "shape"
      ints: 3
      ints: 10
      }
      arg {
      name: "values"
      s:
      "\000\377\152\232\115\072\000\000\200\077\000\377\050\132\215\073\063\063\023\100\000\377\314\063\232\073\000\000\220\100"
      }
     */

    for (auto &o : op.output()) {
      Tensor T;
      if (getConstantByNameOrNull(o)) {
        continue;
      }
      std::vector<dim_t> dim;
      ASSIGN_VALUE_OR_RETURN_ERR(dim, getShape<dim_t>(dict["shape"]));
      T.reset(ElemKind::UInt8QTy, dim, 0.0, 0);
      auto TH = T.getHandle<uint8_t>();
      RETURN_ERR_IF_NOT(
          dict["values"]->strings().size() == 1,
          "Expect single string input for GivenTensorByteStringToUInt8Fill");
      const std::string &str = dict["values"]->strings().Get(0);

      size_t pos;
      for (pos = 0; pos < str.size(); pos++) {
        TH.raw(pos) = (uint8_t)str[pos];
      }

      RETURN_ERR_IF_NOT(
          pos == T.size(),
          strFormat("The number of serialized values (%li) does not "
                    "match the size of the tensor (%li).",
                    pos, (size_t)T.size()));
      RETURN_IF_ERR(createAndRegisterConstant(o, std::move(T)));
    }
    return Error::success();
  }

  // Load quantized tensors:
  if (typeName == "Int8GivenTensorFill" ||
      typeName == "Int8GivenIntTensorFill") {
    /*
     output: "conv1_w"
     name: ""
     type: "Int8GivenTensorFill"
     arg {
     name: "shape"
     ints: 96
     ints: 3
     ints: 11
     ints: 11
     }
     arg {
     name: "values"
     s: "\x7f\x80\x80\x7"
     }
     arg {
     name: "Y_scale"
     f: 0.00044428
     }
     arg {
     name: "Y_zero_point"
     i: 127
     }
     */
    for (auto &o : op.output()) {
      Tensor T;
      if (getConstantByNameOrNull(o)) {
        continue;
      }

      std::vector<dim_t> dim;
      ASSIGN_VALUE_OR_RETURN_ERR(dim, getShape<dim_t>(dict["shape"]));

      RETURN_ERR_IF_NOT(dict.count("Y_zero_point"),
                        ("missing zero point for quantized output type"));
      RETURN_ERR_IF_NOT(dict.count("Y_scale"),
                        ("missing Y_scale for quantized output type"));

      float scale;
      ASSIGN_VALUE_OR_RETURN_ERR(scale, loadFloat(dict["Y_scale"]));
      (void)scale;
      int32_t offset;
      ASSIGN_VALUE_OR_RETURN_ERR(offset, loadInt(dict["Y_zero_point"]));
      (void)offset;
      size_t i = 0;
      if (typeName == "Int8GivenTensorFill") {
        // Although in Caffe2 quantized model, the weights is int8 quantized,
        // the weights is stored in uint8_t format due to that Caffe2 requires
        // the type of input and weights must be the same. Therefore, we need
        // to convert it to int8 by subtracting 128.
        TypeRef ty;
        ASSIGN_VALUE_OR_RETURN_ERR(
            ty, loadQuantTy(o, ElemKind::Int8QTy, dim, dict,
                            /* skipClipQuantRangeToFP16 */ true));
        T.reset(*ty);
        auto TH = T.getHandle<int8_t>();
        std::string str = dict["values"]->s();
        for (; i < str.size(); i++) {
          TH.raw(i) = ((uint8_t)(str.c_str()[i]) - UINT8_TO_INT8_SHIFT);
        }
      } else {
        TypeRef ty;
        ASSIGN_VALUE_OR_RETURN_ERR(
            ty, loadQuantTy(o, ElemKind::Int32QTy, dim, dict,
                            /* skipClipQuantRangeToFP16 */ true));
        T.reset(*ty);
        auto TH = T.getHandle<int32_t>();
        for (auto num : dict["values"]->ints()) {
          TH.raw(i++) = num;
        }
      }

      // If we're clipping quantized ranges tp FP16, then we need to rescale the
      // Tensor and update its type.
      if (clipQuantRangeToFP16_) {
        const ElemKind k = T.getType().getElementType();
        const auto qMinMax = getQuantizedValueRange(T.getType().getScale(),
                                                    T.getType().getOffset(), k);
        const float newMin = std::max(qMinMax.first, kMinFP16);
        const float newMax = std::min(qMinMax.second, kMaxFP16);
        if (newMin != qMinMax.first || newMax != qMinMax.second) {
          auto rescaledT = glow::make_unique<Tensor>();
          dispatchQuantizedImpl(rescaleQTensor, k, T, *rescaledT, newMin,
                                newMax);
          T = std::move(*rescaledT);
        }
      }

      RETURN_ERR_IF_NOT(
          i == T.size(),
          strFormat("The number of serialized values (%li) does not "
                    "match the size of the tensor (%li).",
                    i, (size_t)T.size()));

      RETURN_IF_ERR(createAndRegisterConstant(o, std::move(T)));
    }

    return Error::success();
  }

  // Load tensors with constant fill:
  if (typeName == "ConstantFill") {
    /*
     output: "data"
     name: ""
     type: "ConstantFill"
     arg {
     name: "shape"
     ints: 1
     }
     */

    const auto &name = op.output(0);
    // If the tensor is pre-populated by the user of this class then we don't
    // need to allocate a new tensor.
    if (getConstantByNameOrNull(name)) {
      return Error::success();
    }

    // The shape is set either the shape argument, or from another input
    // tensor. Shape takes priority over input.
    std::vector<dim_t> dims;
    if (dict.count("shape")) {
      ASSIGN_VALUE_OR_RETURN_ERR(dims, getShape<dim_t>(dict["shape"]));
    } else {
      RETURN_ERR_IF_NOT(op.input_size() > 0,
                        "If no shape provided, must have input shape.");

      bool inputAsShape = false;
      if (dict.count("input_as_shape")) {
        ASSIGN_VALUE_OR_RETURN_ERR(inputAsShape,
                                   loadInt(dict["input_as_shape"]));
      }

      if (inputAsShape) {
        // It must be registered as a Constant because it must be statically set
        // already, as shapes must be statically known.
        Constant *in;
        ASSIGN_VALUE_OR_RETURN_ERR(in, getConstantByName(op.input(0)));
        RETURN_ERR_IF_NOT(in->dims().size() == 1,
                          opErrMsg(op, "Input must be 1D tensor."));
        RETURN_ERR_IF_NOT(in->getElementType() == ElemKind::Int64ITy,
                          opErrMsg(op, "Input must be of int64 type."));
        const auto handle = in->getHandle<int64_t>();
        dims.reserve(in->dims().size());
        for (auto dim : handle) {
          dims.push_back(dim);
        }
      } else {
        NodeValue in;
        ASSIGN_VALUE_OR_RETURN_ERR(in, getNodeValueByName(op.input(0)));
        dims = in.dims();
      }
    }

    int to = caffe2::TensorProto_DataType_FLOAT;
    if (dict.count("dtype")) {
      ASSIGN_VALUE_OR_RETURN_ERR(to, loadInt(dict["dtype"]));
    }

    SplatNode *splatNode{nullptr};

    switch (to) {
    case caffe2::TensorProto_DataType_FLOAT: {
      float f = 0.0f;
      if ((dict.count("value") && dict["value"]->has_f())) {
        ASSIGN_VALUE_OR_RETURN_ERR(f, loadFloat(dict["value"]));
      }
      splatNode =
          G_->createSplat(opName, mod_.uniqueType(ElemKind::FloatTy, dims), f);
      break;
    }
    case caffe2::TensorProto_DataType_INT32: {
      int i = 0;
      if ((dict.count("value") && dict["value"]->has_i())) {
        ASSIGN_VALUE_OR_RETURN_ERR(i, loadInt(dict["value"]));
      }
      splatNode =
          G_->createSplat(opName, mod_.uniqueType(ElemKind::Int32ITy, dims), i);
      break;
    }
    case caffe2::TensorProto_DataType_INT64:
    case caffe2::TensorProto_DataType_BOOL: {
      int i = 0;
      if ((dict.count("value") && dict["value"]->has_i())) {
        ASSIGN_VALUE_OR_RETURN_ERR(i, loadInt(dict["value"]));
      }
      splatNode =
          G_->createSplat(opName, mod_.uniqueType(ElemKind::Int64ITy, dims), i);
      break;
    }
    default:
      return MAKE_ERR("Unsupported datatype for ConstantFill.");
    }

    RETURN_IF_ERR(addNodeAsOutput(op, splatNode));

    return Error::success();
  }

  if (typeName == "UniformFill") {
    /*
     output: "fc/w"
     name: ""
     type: "UniformFill"
     arg {
       name: "max"
       f: 0.25
     }
     arg {
       name: "shape"
       ints: 1
       ints: 16
     }
     arg {
       name: "min"
       f: -0.25
     }
    */
    const auto &name = op.output(0);
    Tensor T;
    std::vector<dim_t> dim;
    if (dict.count("shape")) {
      ASSIGN_VALUE_OR_RETURN_ERR(dim, getShape<dim_t>(dict["shape"]));
    } else {
      RETURN_ERR_IF_NOT(op.input_size() > 0,
                        "If no shape provided, must have input shape.");

      bool inputAsShape = false;
      if (dict.count("input_as_shape")) {
        ASSIGN_VALUE_OR_RETURN_ERR(inputAsShape,
                                   loadInt(dict["input_as_shape"]));
      }

      if (inputAsShape) {
        Constant *in;
        ASSIGN_VALUE_OR_RETURN_ERR(in, getConstantByName(op.input(0)));
        RETURN_ERR_IF_NOT(in->dims().size() == 1,
                          opErrMsg(op, "Input must be 1D tensor."));
        RETURN_ERR_IF_NOT(in->getElementType() == ElemKind::Int64ITy,
                          opErrMsg(op, "Input must be of int64 type."));
        const auto handle = in->getHandle<int64_t>();
        dim.reserve(in->dims().size());
        for (auto d : handle) {
          dim.push_back(d);
        }
      } else {
        NodeValue input;
        ASSIGN_VALUE_OR_RETURN_ERR(input, getNodeValueByName(op.input(0)));
        dim = input.dims();
      }
    }
    T.reset(ElemKind::FloatTy, dim);
    auto TH = T.getHandle<>();
    float tensorMin;
    ASSIGN_VALUE_OR_RETURN_ERR(tensorMin, loadFloat(dict["min"]));
    float tensorMax;
    ASSIGN_VALUE_OR_RETURN_ERR(tensorMax, loadFloat(dict["max"]));

    DLOG(INFO)
        << "The model contains UniformFill operator, which generates random "
           "numbers. This could be source of discrepancy.";

    // Uniformly generate random numbers in [tensorMin; tensorMax).
    for (auto &elem : TH) {
      elem = mod_.getPRNG().nextRandReal(tensorMin, tensorMax);
    }

    RETURN_IF_ERR(createAndRegisterConstant(name, std::move(T)));

    return Error::success();
  }

  // Load tensors with constant fill:
  if (typeName == "GaussianFill") {
    /*
     output: "data"
     name: ""
     type: "GaussianFill"
     arg {
       name: "mean"
       f: 0.0
     }
     arg {
       name: "std"
       f: 1.0
     }
     arg {
       name: "shape"
       ints: 1
       ints: 16
     }
     */

    const auto &name = op.output(0);
    if (getConstantByNameOrNull(name)) {
      return Error::success();
    }

    // The shape of the output is set by shape, if provided. Otherwise, it is
    // set by the shape of the input or the shape indicated by input if
    // input_as_shape is true
    NodeValue input;
    std::vector<dim_t> dims;
    if (dict.count("shape")) {
      ASSIGN_VALUE_OR_RETURN_ERR(dims, getShape<dim_t>(dict["shape"]));
    } else {
      RETURN_ERR_IF_NOT(op.input_size() > 0,
                        "If no shape provided, must have input shape.");

      bool inputAsShape = false;
      if (dict.count("input_as_shape")) {
        ASSIGN_VALUE_OR_RETURN_ERR(inputAsShape,
                                   loadInt(dict["input_as_shape"]));
      }

      if (inputAsShape) {
        Constant *in;
        ASSIGN_VALUE_OR_RETURN_ERR(in, getConstantByName(op.input(0)));
        RETURN_ERR_IF_NOT(in->dims().size() == 1,
                          opErrMsg(op, "Input must be 1D tensor."));
        RETURN_ERR_IF_NOT(in->getElementType() == ElemKind::Int64ITy,
                          opErrMsg(op, "Input must be of int64 type."));
        const auto handle = in->getHandle<int64_t>();
        dims.reserve(in->dims().size());
        for (auto dim : handle) {
          dims.push_back(dim);
        }
      } else {
        ASSIGN_VALUE_OR_RETURN_ERR(input, getNodeValueByName(op.input(0)));
        dims = input.dims();
      }

      if (dict.count("extra_shape")) {
        std::vector<dim_t> extra_shape;
        ASSIGN_VALUE_OR_RETURN_ERR(extra_shape,
                                   getShape<dim_t>(dict["extra_shape"]));
        dims.insert(dims.end(), extra_shape.begin(), extra_shape.end());
      }
    }
    if ((!input && !dims.empty()) || input.dims().vec() != dims) {
      input =
          G_->createSplat("in", mod_.uniqueType(ElemKind::FloatTy, dims), 0.);
    }
    float mean;
    ASSIGN_VALUE_OR_RETURN_ERR(mean, loadFloat(dict["mean"]));
    float scale;
    ASSIGN_VALUE_OR_RETURN_ERR(scale, loadFloat(dict["std"]));

    auto GF = G_->createGaussianFill(opName, input, mean, scale,
                                     std::random_device{}());
    auto outputType =
        mod_.uniqueType(ElemKind::FloatTy, GF->getResult().dims());
    auto node = G_->createConvertTo(opName + ".ConvertOutput", GF, outputType);
    RETURN_IF_ERR(addNodeAsOutput(op, node));

    return Error::success();
  }

  return MAKE_ERR(unexpectedNodeErrorMessage(op, "Unsupported weight kind"));
}