TensorRT 10.14+ 多 GPU 构建回归与 Cask 支持的运算

def make_model() -> bytes: x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 1, 8, 8]) node = helper.make_node( "MaxPool", ["x"], ["y"], name="MaxPool_0", kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1], ) y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 1, 8, 8]) graph = helper.make_graph([node], "tiny_maxpool", [x], [y]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 20)]) model.ir_version = 9 onnx.checker.check_model(model) return model.SerializeToString() def build(model: bytes, cudart: ctypes.CDLL, device: int, label: str) -> bool: assert cudart.cudaSetDevice(device) == 0 print(f"======={label} on {device=}", file=sys.stderr) logger = trt.Logger(trt.Logger.INFO) builder = trt.Builder(logger) config = builder.create_builder_config() config.set_memory_pool_limit( trt.MemoryPoolType.WORKSPACE, 1 << 30 ) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED) ) parser = trt.OnnxParser(network, logger) assert parser.parse(model, path="/tmp") engine = builder.build_serialized_network(network, config) print("=======Failed" if engine is None else "Succeed", file=sys.stderr) def main(): print(f"======={trt.__version__=}", file=sys.stderr) cudart = ctypes.CDLL("libcudart.so") count = ctypes.c_int() cudart.cudaGetDeviceCount(ctypes.byref(count)) print(f"======={count.value=}", file=sys.stderr) if count.value < 2: return 1 model = make_model() build(model, cudart, 0, "same device 1/2") build(model, cudart, 0, "same device 2/2") # succeeds build(model, cudart, 0, "different device 1/2") build(model, cudart, 1, "different device 2/2") # fails