MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / dump_tensor

Method dump_tensor

src/serialization/impl/serializer_oss.cpp:406–465  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

404}
405
406void GraphDumperOSS::dump_tensor(
407 const std::string& name, const HostTensorND& tensor, TensorWriteMethod method,
408 TensorFormat) {
409 using namespace flatbuffers;
410 using Meth = TensorWriteMethod;
411 mgb_assert(
412 (method == Meth::VALUE_ANONYMOUS) ^ (!name.empty()),
413 "name must be non-empty for non Meth::VALUE_ANONYMOUS tensors");
414
415 bool has_value = method != Meth::META_INPUT;
416 bool should_keep_name = true;
417 switch (method) {
418 case Meth::VALUE_ANONYMOUS:
419 should_keep_name = false;
420 break;
421 case Meth::VALUE_SHARED:
422 should_keep_name = m_config.keep_param_name;
423 ++m_nr_shared_tensor;
424 if (m_config.keep_param_name) {
425 mgb_assert(
426 m_used_param_names.insert(name).second,
427 "duplicated VALUE_SHARED tensor name: %s", name.c_str());
428 m_cur_rst.params.emplace_back(name);
429 }
430 break;
431 case Meth::META_INPUT:
432 case Meth::VALUE_INPUT:
433 mgb_assert(!name.empty(), "empty input tensor name");
434 mgb_assert(
435 m_used_input_names.insert(name).second,
436 "duplicated input tensor name: %s", name.c_str());
437 m_cur_rst.inputs.emplace_back(name);
438 break;
439 }
440
441 size_t value_size = 0;
442 if (has_value) {
443 check_tensor_value_valid(name, tensor);
444 auto begin = m_file->tell();
445 auto&& dumper = m_config.tensor_value_dumper;
446 if (dumper) {
447 dumper(*m_file, *m_cur_opr, tensor);
448 } else {
449 m_file->write(tensor.raw_ptr(), tensor.layout().span().high_byte);
450 }
451 value_size = m_file->tell() - begin;
452 m_cur_rst.tensor_value_bytes += value_size;
453 }
454
455 auto fbname = should_keep_name ? m_builder.CreateSharedString(name) : 0;
456 auto shape = m_builder.CreateVectorScalarCast<uint32_t>(
457 tensor.shape().shape, tensor.shape().ndim);
458 auto comp_node = fbs::CreateCompNode(
459 m_builder,
460 m_builder.CreateSharedString(tensor.comp_node().to_string_logical()));
461 auto dtype = build_dtype(tensor.dtype());
462 auto serialized_tensor =
463 fbs::CreateTensor(m_builder, fbname, shape, comp_node, dtype, value_size);

Callers

nothing calls this directly

Calls 15

check_tensor_value_validFunction · 0.85
build_dtypeFunction · 0.85
CreateTensorClass · 0.85
emplace_backMethod · 0.80
spanMethod · 0.80
to_string_logicalMethod · 0.80
emptyMethod · 0.45
insertMethod · 0.45
tellMethod · 0.45
writeMethod · 0.45
raw_ptrMethod · 0.45
layoutMethod · 0.45

Tested by

no test coverage detected