экспорт pytorch в onnx с определенными выходными слоями

Question

Я пытаюсь экспортировать сеть pytorch в ONNX, а затем onnx в tenorrt для более быстрого вывода, но мне нужно получить вывод слоя (avg_pool) до окончательного слоя softmax или любого промежуточного слоя, модели [-N], так что согласно следующему image слой avgpool называется GlobalAveragePool, который возвращает 512- мерный вектор признаков в качестве выходного тензора, здесь я использую предварительно обученную модель resnet-18 в качестве образца сетевой магистрали.

поэтому мой запрос заключается в том, как это сделать с помощью сценария экспорта onnx, ниже приведен код, который определяет привязки ввода, вывода и экспортирует модель,

           model = torchvision.models.resnet18(pretrained=True)
    device = torch.device('cuda:0')
    model.to(device).eval()
    example = torch.rand(1, 3, 128, 128).to(device)
    f = 'saved_weight/weight.onnx'  # filename
    torch.onnx.export(model, example, f, verbose=False, opset_version=12, input_names=['images'],
                              output_names=['output'])

Но я не могу понять, как указать выходной слой во время экспорта модели, а также следующее, как я указываю ввод, привязку вывода в тензоре, чтобы вывести модель и получить выходные тензоры,

       engine, context = build_engine(PLAN_FILE_PATH)
# get sizes of input and output and allocate memory required for input data and for output data
    for binding in engine:
        print(binding)
        if engine.binding_is_input(binding):  # we expect only one input
            input_shape = engine.get_binding_shape(binding)
            # print(input_shape)
            input_size = trt.volume(
                input_shape) * engine.max_batch_size * np.dtype(np.float32).itemsize  # in bytes
            device_input = cuda.mem_alloc(input_size)
        else:  # and one output
            output_shape = engine.get_binding_shape(binding)
            # print(output_shape)
            # create page-locked memory buffers (i.e. won't be swapped to disk)
            host_output = cuda.pagelocked_empty(trt.volume(
                output_shape) * engine.max_batch_size, dtype=np.float32)
            device_output = cuda.mem_alloc(host_output.nbytes)

# preprocess input data
host_input = np.array(preprocess_image(img_path),dtype=np.float32, order='C')
print(host_input.shape)
cuda.memcpy_htod_async(device_input, host_input, stream)

# run inference
context.execute_async(bindings=[int(device_input), int(device_output)],stream_handle=stream.handle)

Таким образом, во время вывода с использованием tenorrt выходные привязки будут теми, которые были указаны во время экспорта модели, поэтому в случае наличия нескольких выходных данных мы можем просто указать выходную привязку, чтобы получить тензоры для указанного слоя.