Program Listing for File add_classification.cpp

Return to documentation for file (morpheus/_lib/src/stages/add_classification.cpp)

Copy
Copied!
            

#include "morpheus/stages/add_classification.hpp" #include "morpheus/objects/dev_mem_info.hpp"// for DevMemInfo #include "morpheus/objects/tensor.hpp" #include "morpheus/objects/tensor_object.hpp"// for TensorIndex, TensorObject #include "morpheus/utilities/matx_util.hpp" #include "morpheus/utilities/tensor_util.hpp"// for TensorUtils::get_element_stride #include "morpheus/utilities/type_util.hpp"// for DType #include "morpheus/utilities/type_util_detail.hpp"// for DataType #include <cuda_runtime.h>// for cudaMemcpy, cudaMemcpyDeviceToDevice #include <glog/logging.h> #include <mrc/cuda/common.hpp>// for MRC_CHECK_CUDA #include <rmm/cuda_stream_view.hpp>// for cuda_stream_per_thread #include <rmm/device_buffer.hpp>// for device_buffer #include <algorithm>// for min_element, transform #include <cstddef> #include <exception> #include <functional>// for divides, bind, placeholders #include <memory> #include <ostream>// needed for logging #include <type_traits>// for declval #include <utility>// for move // IWYU thinks we need __alloc_traits<>::value_type for vector assignments // IWYU pragma: no_include <ext/alloc_traits.h> namespace morpheus { // Component public implementations // ************ AddClassificationStage **************************** // AddClassificationsStage::AddClassificationsStage(float threshold, std::size_t num_class_labels, std::map<std::size_t, std::string> idx2label) : PythonNode(base_t::op_factory_from_sub_fn(build_operator())), m_threshold(threshold), m_num_class_labels(num_class_labels), m_idx2label(std::move(idx2label)) { CHECK(m_idx2label.size() <= m_num_class_labels) << "idx2label should represent a subset of the class_labels"; } AddClassificationsStage::subscribe_fn_t AddClassificationsStage::build_operator() { return [this](rxcpp::observable<sink_type_t> input, rxcpp::subscriber<source_type_t> output) { return input.subscribe(rxcpp::make_observer<sink_type_t>( [this, &output](sink_type_t x) { const auto& probs = x->get_probs(); const auto& shape = probs.get_shape(); // Depending on the input the stride is given in bytes or elements, convert to elements auto stride = TensorUtils::get_element_stride<std::size_t>(probs.get_stride()); CHECK(shape.size() == 2 && shape[1] == m_num_class_labels) << "Label count does not match output of model. Label count: " << m_num_class_labels << ", Model output: " << shape[1]; const std::size_t num_rows = shape[0]; const std::size_t num_columns = shape[1]; // A bit ugly, but we cant get access to the rmm::device_buffer here. So make a copy auto tmp_buffer = std::make_shared<rmm::device_buffer>(probs.bytes(), rmm::cuda_stream_per_thread); MRC_CHECK_CUDA( cudaMemcpy(tmp_buffer->data(), probs.data(), tmp_buffer->size(), cudaMemcpyDeviceToDevice)); // Now call the threshold function auto thresh_bool_buffer = MatxUtil::threshold(DevMemInfo{tmp_buffer, probs.dtype(), shape, stride}, m_threshold, false); std::vector<TensorIndex> tensor_shape(shape.size()); std::copy(shape.cbegin(), shape.cend(), tensor_shape.begin()); std::vector<TensorIndex> tensor_stride(stride.size()); std::copy(stride.cbegin(), stride.cend(), tensor_stride.begin()); auto tensor_obj = Tensor::create(thresh_bool_buffer, DType::create<bool>(), tensor_shape, tensor_stride); std::vector<std::string> columns(m_idx2label.size()); std::vector<TensorObject> tensors(m_idx2label.size()); std::size_t i = 0; for (const auto& [column_num, column_name] : m_idx2label) { columns[i] = column_name; tensors[i] = tensor_obj.slice(std::vector<TensorIndex>{0, static_cast<TensorIndex>(column_num)}, std::vector<TensorIndex>{static_cast<TensorIndex>(num_rows), static_cast<TensorIndex>(column_num + 1)}); ++i; } x->set_meta(columns, tensors); output.on_next(x); }, [&](std::exception_ptr error_ptr) { output.on_error(error_ptr); }, [&]() { output.on_completed(); })); }; } // ************ AddClassificationStageInterfaceProxy ************* // std::shared_ptr<mrc::segment::Object<AddClassificationsStage>> AddClassificationStageInterfaceProxy::init( mrc::segment::Builder& builder, const std::string& name, float threshold, std::size_t num_class_labels, std::map<std::size_t, std::string> idx2label) { auto stage = builder.construct_object<AddClassificationsStage>(name, threshold, num_class_labels, idx2label); return stage; } } // namespace morpheus

© Copyright 2023, NVIDIA. Last updated on Feb 3, 2023.