| 94 | } |
| 95 | |
| 96 | void TensorRTRuntimeOpr::get_output_var_shape( |
| 97 | const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const { |
| 98 | auto batch = inp_shape.at(0)[0]; |
| 99 | m_manager.create_trt_context(this->comp_node(), inp_shape, m_engine.get()); |
| 100 | auto get_mgb_shape = [&](int binding_idx) -> TensorShape { |
| 101 | auto dims = m_engine->getBindingDimensions(binding_idx); |
| 102 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 103 | auto format = m_engine->getBindingFormat(binding_idx); |
| 104 | // converting dims to nchw4 format |
| 105 | if (format == nvinfer1::TensorFormat::kCHW4) { |
| 106 | mgb_assert( |
| 107 | dims.nbDims == 3 || dims.nbDims == 4, |
| 108 | "Tensor with NCHW4 format should have dimensions of " |
| 109 | "3/4.(got: %d)", |
| 110 | dims.nbDims); |
| 111 | int chan_pos = 0; |
| 112 | if (dims.nbDims == 4) { |
| 113 | chan_pos = 1; |
| 114 | } |
| 115 | dims.nbDims = dims.nbDims + 1; |
| 116 | dims.d[chan_pos] = (dims.d[chan_pos] + 3) / 4; |
| 117 | dims.d[dims.nbDims - 1] = 4; |
| 118 | } |
| 119 | #endif |
| 120 | auto shape = m_trt_engine_has_batch ? TensorRTOpr::dims2shape(dims) |
| 121 | : TensorRTOpr::dims2shape(dims, batch); |
| 122 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 123 | if (static_cast<size_t>(binding_idx) < inp_shape.size()) { |
| 124 | for (int i = 0; i < dims.nbDims; i++) { |
| 125 | if (dims.d[i] == -1) { |
| 126 | shape[i] = inp_shape.at(binding_idx)[i]; |
| 127 | } |
| 128 | } |
| 129 | } else { |
| 130 | auto trt_infer_dims = m_manager.get_binding_dimensions(binding_idx); |
| 131 | for (int i = 0; i < dims.nbDims; i++) { |
| 132 | if (dims.d[i] == -1) { |
| 133 | shape[i] = trt_infer_dims.d[i]; |
| 134 | } |
| 135 | } |
| 136 | } |
| 137 | #endif |
| 138 | return shape; |
| 139 | }; |
| 140 | for (size_t i = 0; i < inp_shape.size(); ++i) { |
| 141 | mgb_assert(batch == inp_shape[i][0], "input batchsize not equal"); |
| 142 | TensorShape shp = get_mgb_shape(i); |
| 143 | mgb_assert( |
| 144 | shp.eq_shape(inp_shape[i]), "input shape mismatch: expect=%s got=%s", |
| 145 | shp.to_string().c_str(), inp_shape[i].to_string().c_str()); |
| 146 | } |
| 147 | for (size_t i = 0; i < out_shape.size() - 1; ++i) { |
| 148 | out_shape[i] = get_mgb_shape(i + input().size()); |
| 149 | } |
| 150 | out_shape.back() = {intl::workspace_size(m_engine.get())}; |
| 151 | } |
| 152 | |
| 153 | void TensorRTRuntimeOpr::add_input_layout_constraint() { |
nothing calls this directly
no test coverage detected