Method set_inputs

monai/networks/trt_compiler.py:168–202

Sets input bindings for TRT engine according to feed_dict Args: feed_dict: a dictionary [str->Tensor] stream: CUDA stream to use

(self, feed_dict, stream)

                ctx.set_tensor_address(binding, t.data_ptr())

    def set_inputs(self, feed_dict, stream):
        """
        Sets input bindings for TRT engine according to feed_dict
        Args:
           feed_dict: a dictionary [str->Tensor]
           stream: CUDA stream to use
        """
        e = self.engine
        ctx = self.context

        last_profile = self.cur_profile

        def try_set_inputs():
            for binding in self.input_names:
                t = feed_dict.get(self.input_table[binding], None)
                if t is not None:
                    t = t.contiguous()
                    shape = t.shape
                    ctx.set_input_shape(binding, shape)
                    ctx.set_tensor_address(binding, t.data_ptr())

        while True:
            try:
                try_set_inputs()
                break
            except ShapeError:
                next_profile = (self.cur_profile + 1) % e.num_optimization_profiles
                if next_profile == last_profile:
                    raise
                self.cur_profile = next_profile
                ctx.set_optimization_profile_async(self.cur_profile, stream)
            except Exception:
                raise
        left = ctx.infer_shapes()
        assert len(left) == 0

    def infer(self, stream, use_cuda_graph=False):
        """