| 570 | } |
| 571 | |
| 572 | void TensorRTOpr::get_output_var_shape( |
| 573 | const TensorShapeArray& inp_shape, TensorShapeArray& out_shape) const { |
| 574 | for (size_t i = 0; i < inp_shape.size(); ++i) { |
| 575 | set_input_by_tensor_shape(m_network->getInput(i), inp_shape[i]); |
| 576 | } |
| 577 | |
| 578 | for (int i = 0; i < m_network->getNbOutputs(); ++i) { |
| 579 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 580 | auto output = m_network->getOutput(i); |
| 581 | out_shape[i] = dims2shape(output->getDimensions()); |
| 582 | auto tensor_format = output->getAllowedFormats(); |
| 583 | // fix tensor shape from tensor format |
| 584 | if (tensor_format & (1 << static_cast<int>(nvinfer1::TensorFormat::kCHW4))) { |
| 585 | mgb_assert(out_shape[i].ndim == 4); |
| 586 | out_shape[i].ndim++; |
| 587 | out_shape[i].shape[1] /= 4; |
| 588 | out_shape[i].shape[4] = 4; |
| 589 | } |
| 590 | #else |
| 591 | out_shape[i] = dims2shape(m_network->getOutput(i)->getDimensions()); |
| 592 | #endif |
| 593 | } |
| 594 | |
| 595 | // Because input shape is NCHW, so the batch size should always be 1. |
| 596 | m_builder->setMaxBatchSize(1); |
| 597 | |
| 598 | auto self = const_cast<TensorRTOpr*>(this); |
| 599 | if (m_engine == nullptr && TensorRTEngineCache::enable_engine_cache()) { |
| 600 | self->build_engine_from_cache(); |
| 601 | } |
| 602 | |
| 603 | bool engine_valid = true; |
| 604 | if (m_engine == nullptr) { |
| 605 | engine_valid = false; |
| 606 | } else { |
| 607 | int nr_input = m_network->getNbInputs(); |
| 608 | mgb_assert( |
| 609 | static_cast<size_t>(nr_input) == input().size(), "input size changed"); |
| 610 | for (int i = 0; i < nr_input; ++i) { |
| 611 | int binding_idx = |
| 612 | m_engine->getBindingIndex(m_network->getInput(i)->getName()); |
| 613 | auto cuda_engine_shp = |
| 614 | dims2shape(m_engine->getBindingDimensions(binding_idx)); |
| 615 | #if NV_TENSOR_RT_VERSION >= 6001 |
| 616 | auto tensor_format = m_engine->getBindingFormat(binding_idx); |
| 617 | // fix tensor shape from tensor format |
| 618 | if (tensor_format == nvinfer1::TensorFormat::kCHW4) { |
| 619 | mgb_assert(cuda_engine_shp.ndim == 4); |
| 620 | cuda_engine_shp.ndim++; |
| 621 | cuda_engine_shp[1] /= 4; |
| 622 | cuda_engine_shp[4] = 4; |
| 623 | } |
| 624 | #endif |
| 625 | if (!cuda_engine_shp.eq_shape(inp_shape[i])) { |
| 626 | engine_valid = false; |
| 627 | break; |
| 628 | } |
| 629 | } |
nothing calls this directly
no test coverage detected