Function UndoWeightsShuffling

tensorflow/lite/toco/tooling_util.cc:2387–2432

}

void UndoWeightsShuffling(Model* model) {
  for (const auto& op : model->operators) {
    if (op->type != toco::OperatorType::kFullyConnected) {
      continue;
    }
    const auto& fc_op = static_cast<toco::FullyConnectedOperator&>(*op);
    if (fc_op.weights_format == FullyConnectedWeightsFormat::kDefault) {
      continue;
    }
    const string& weights_name = fc_op.inputs[1];
    QCHECK_EQ(CountOpsWithInput(*model, weights_name), 1);
    auto& weights_array = model->GetArray(weights_name);
    QCHECK(weights_array.data_type == ArrayDataType::kUint8);
    auto& weights_data =
        weights_array.GetMutableBuffer<toco::ArrayDataType::kUint8>().data;
    const auto& weights_shape = weights_array.shape();
    QCHECK_EQ(weights_shape.dimensions_count(), 2);
    const int rows = weights_shape.dims(0);
    const int cols = weights_shape.dims(1);
    QCHECK_EQ(rows % 4, 0);
    QCHECK_EQ(cols % 16, 0);
    CHECK_EQ(rows * cols, weights_data.size());
    // Compute the de-shuffled weights
    std::vector<uint8> deshuffled_data(weights_data.size());
    uint8* shuffled_data_ptr = weights_data.data();
    for (int r = 0; r < rows; r += 4) {
      for (int c = 0; c < cols; c += 16) {
        for (int i = 0; i < 4; i++) {
          uint8* deshuffled_data_ptr =
              deshuffled_data.data() + (r + i) * cols + c;
          for (int j = 0; j < 16; j++) {
            uint8 shuffled_val = *shuffled_data_ptr++;
            // Deshuffling isn't only about deshuffling the storage layout,
            // it's also about undoing the flipping of the sign bit, which is
            // performed on the shuffled weights.
            uint8 deshuffled_val = shuffled_val ^ 0x80;
            *deshuffled_data_ptr++ = deshuffled_val;
          }
        }
      }
    }
    CHECK_EQ(shuffled_data_ptr, weights_data.data() + rows * cols);
    // Switch this FC op to using the deshuffled weights.
    weights_data = std::move(deshuffled_data);
  }
}

void CopyMinMaxAndQuantizationRelatedFields(const Array& src, Array* dst) {
  if (src.minmax) {