multi-context_inference.py

# Create multiple optimization profiles for different contexts to use
shape0 = (1, 3, 224, 224)
profile0 = builder.create_optimization_profile()
profile0.set_shape("input", min=shape0, opt=shape0, max=shape0)
config.add_optimization_profile(profile0)

shape1 = (1, 3, 448, 448)
profile1 = builder.create_optimization_profile()
profile1.set_shape("input", min=shape1, opt=shape1, max=shape1)
config.add_optimization_profile(profile1)

# Create two contexts, each set to a different optimization profile
context0 = engine.create_execution_context()
context0.active_optimization_profile = 0
binding_index0 = engine.get_binding_index("input")
context0.set_binding_dimensions(binding_index0, shape0)

context1 = engine.create_execution_context()
context1.active_optimization_profile = 1
binding_index1 = engine.get_binding_index("input [profile 1]")
context1.set_binding_dimensions(binding_index1, shape1)

# Assuming the output bindings/buffers are large enough 
# to hold the outputs of either input shape. Alternatively,
# you could create separate bindings0 and bindings1 instead.
if input_shape == shape0:
  context0.execute_v2(bindings)
elif input_shape == shape1:
  context1.execute_v2(bindings)