[ONNX] ONNX 배치 사이즈 변경하는 방법 + 삽질

 

onnx batch size change 라는 키워드를 통해 나온 검색 결과에 의하면,

다음 코드를 사용하면 기존에 만들어 두었던 onnx 파일의 배치 사이즈를 변경할 수 있다고 하는데, 

아래 코드는 입력 사이즈만 변경하는 코드이다. 절대 아래와 같이 사용하면 안된다. 

import onnx

def change_input_dim(model):
    # Use some symbolic name not used for any other dimension
    sym_batch_dim = "N"
    # or an actal value
    actual_batch_dim = 4 

    # The following code changes the first dimension of every input to be batch-dim
    # Modify as appropriate ... note that this requires all inputs to
    # have the same batch_dim 
    inputs = model.graph.input
    for input in inputs:
        # Checks omitted.This assumes that all inputs are tensors and have a shape with first dim.
        # Add checks as needed.
        dim1 = input.type.tensor_type.shape.dim[0]
        # update dim to be a symbolic value
        dim1.dim_param = sym_batch_dim
        # or update it to be an actual value:
        # dim1.dim_value = actual_batch_dim


def apply(transform, infile, outfile):
    model = onnx.load(infile)
    transform(model)
    onnx.save(model, outfile)

apply(change_input_dim, r"input-file-name, r"output-file-name")

 

이는 입력 배치 사이즈는 변경해주지만, 출력 배치 사이즈는 변경하지 않는다. 실제로 netron 을 사용하여 그래프를 확인해보면 출력 단에서는 배치 사이즈가 변경 되지 않았음을 알 수 있다. 

 

즉, 입출력에 예민한 TensorRT에서 배치 사이즈를 1이 아닌 값으로 정적(no Dynamic)으로 고정하여 엔진을 생성할 때는

위와 같은 방법으로 생성된 onnx 모델을 사용하면 절대 안되므로, 

onnx.export 를 사용하여 원하는 배치 사이즈를 가지는 텐서 값을 통해 정상적으로 onnx 모델을 생성하여 TensorRT을 생성 해야한다. 

저 코드를 알려준 저자에게 -1 눌러주고 왔다....

 

 

관련 이슈 1 : github.com/onnx/keras-onnx/issues/605

Explicitly set batch size of generate ONNX model · Issue #605 · onnx/keras-onnx

관련 이슈 2: github.com/onnx/onnx/issues/2182

Changing Batch SIze · Issue #2182 · onnx/onnx