pytorch inference 총망라

결과

Model	FP32	FP16	INT8(PTQ)
U2Net(torch)	23.18ms	15.34ms
U2Net(torch script)	21.36ms	10.98ms
U2Net(onnx runtime)	17.48ms	12.88ms
U2Net(tensorrt)	13.31ms	6.32ms	5.20ms

[1] pytorch 모델

pytorch 에서는 .pt, .pth 를 사용하길 권장하고 있다.

그러나 여기에서.pth는 Python path (.pth) configuration files 과 혼동될 수 있다고 한다.

이건 중요한게 아니고 pytorch에서는 weights만 저장하는 방식과 model을 통채로 저장하는 방식이 있는데 두 방식 모두 같은 확장자를 사용하는 것 같다.

(개인적인 생각이지만 이게 더 헷갈린다. 찾아봐도 이걸 따로 구분하려는 의지는 커뮤니티에서 없는것 같다.)

당연하겠지만 모델의 구조를 기술하는데 추가적인 크기가 크지 않으므로 두 방식의 파일 크기는 거의 차이가 없다고 봐도 된다.

model = U2NET(3, 1)
model.load_state_dict(torch.load("u2net(original).pth"))
model.cuda().eval()

# 모델 전체 저장(fp32)
torch.save(model, "u2net(saved_model,fp32).pt")

# 모델의 가중치만 저장(fp32)
torch.save(model.state_dict(), "u2net(state_dict,fp32).pt")

model.half()
# 모델 전체 저장(fp16)
torch.save(model, "u2net(saved_model,fp16).pt")

# 모델의 가중치만 저장(fp16)
torch.save(model.state_dict(), "u2net(state_dict,fp16).pt")

fp16 으로 변환후에 저장하면 모델의 크기는 반으로 줄어들고, 해당 모델은 fp16 으로 추론된다.

U2Net은 기본 모델이고 U2NetFix 는 deprecated된 함수를 수정한 모델, U2Net_full은 중복 코드를 일반화한 모델이다.

Model	GPU Time(avg)	miou
U2Net(fp32)	23.11ms	0.8978
U2Net(fp16)	14.72ms	0.8974
U2NetFix(fp32)	23.18ms	0.8755
U2NetFix(fp16)	15.34ms	0.8752
U2Net_full(fp32)	23.04ms	0.8936
U2Net_full(fp16)	15.49ms	0.8946

[2] torch script 모델

ts 변환은 U2Net full은 변환이 되지 않는다.
forward 에서 입력데이터의 크기가 변해서 변환에 오류가 발생하는 것인데, U2Net full에서는 중첩함수를 사용해서 그렇다.

model = U2NETFix(3, 1)
model.load_state_dict(torch.load("u2net(original).pth"))
model.cuda().eval()
# torch script 저장(fp32)
ts_model = torch.jit.script(model)
ts_model.save("u2net(ts,fp32).ts")
model.half()
# torch script 저장(fp16)
ts_model = torch.jit.script(model)
ts_model.save("u2net(ts,fp16).ts")

torch script 모델은 첫 실행에 시간이 비정상으로 오래 걸리는데 원인은 모르겠다. 그래서 elapsed time을 정렬해서 느린 10개를 제외하고 측정하였다.

Model	GPU Time(avg)	miou
U2Net(fp32)(ts)	21.81ms	0.8978
U2Net(fp16)(ts)	11.45ms	0.8974
U2NetFix(fp32)(ts)	21.36ms	0.8755
U2NetFix(fp16)(ts)	10.98ms	0.8752

[3] onnx 모델

pip install onnxruntime-gpu

onnx runtime 모델로의 변환코드는 아래와 같다.

model = U2NETFix(3, 1)
model.load_state_dict(torch.load("u2net(original).pth"))
model.cuda().eval()
dummy_input = torch.randn((1, 3, 320, 320)).cuda()
# onnx fp32 변환
torch.onnx.export(model, dummy_input, "u2net(onnx,fp32).onnx",
                  opset_version=11,
                  input_names=['input'],
                  output_names=['output'],
                  )
model.half()
# onnx fp16 변환
torch.onnx.export(model, dummy_input.half(), "u2net(onnx,fp16).onnx",
                  opset_version=11,
                  input_names=['input'],
                  output_names=['output'],
                  )

onnx runtime 인퍼런스는 아래와 같다.
onnx에서는 fp32,fp16 모두 input은 np.ndarray여야 한다.
출력도 np.ndarray로 나온다.

import onnxruntime as ort
ort_sess = ort.InferenceSession("u2net(onnx,fp32).onnx", providers=['CUDAExecutionProvider', 'TensorrtExecutionProvider', 'CPUExecutionProvider'])
d1, d2, d3, d4, d5, d6, d7 = ort_session.run(
            None,
            {'input': img},
    )

변환 작업에서 경고가 뜨긴 하지만 변환에는 모두 성공하고, 추론도 잘 돌아간다.

Model	GPU Time(avg)	miou
U2Net(fp32)(ort)	17.81ms	0.8978
U2Net(fp16)(ort)	13.07ms	0.8974
U2NetFix(fp32)(ort)	17.48ms	0.8755
U2NetFix(fp16)(ort)	12.88ms	0.8752
U2Net_full(fp32)(ort)	17.14ms	0.8936
U2Net_full(fp16)(ort)	12.95ms	0.8946

아래는 provider에 따른 시간 비교이다. (U2NetFix 사용)

Provider	fp32 time	fp16 time
TensorrtExecutionProvider	13.84ms	12.93ms
CUDAExecutionProvider	16.81ms	12.32ms
CPUExecutionProvider	132.98ms	278.30ms
TensorrtExecutionProvider, CUDAExecutionProvider	13.76ms	12.86ms
TensorrtExecutionProvider, CPUExecutionProvider	13.85ms	12.99ms
CUDAExecutionProvider, CPUExecutionProvider	16.94ms	12.21ms
TensorrtExecutionProvider, CUDAExecutionProvider, CPUExecutionProvider	13.77ms	12.94ms

[4] tensorrt 모델

tensorrt 모델은 trtexec로 변환한다.

trtexec --onnx=<ONNX모델경로> --workspace=2048 <--fp16,--int8,--best> --explicitBatch --saveEngine=<TRT출력경로> 

workspace는 cudnn forward에 필요한 cudnnGetConvolutionForwardWorkspaceSize를 지정하는것 같은데 2048정도면 충분하다.

trtexec --onnx=u2net(onnx,fp32).onnx --workspace=2048 --noTF32 --saveEngine=u2net(fp32_notf32).trt
trtexec --onnx=u2net(onnx,fp32).onnx --workspace=2048 --saveEngine=u2net(fp32_fp32).trt
trtexec --onnx=u2net(onnx,fp32).onnx --workspace=2048 --fp16 --saveEngine=u2net(fp32_fp16).trt
trtexec --onnx=u2net(onnx,fp32).onnx --workspace=2048 --int8 --saveEngine=u2net(fp32_int8).trt

Model	GPU Time(avg)	miou	Throughput
U2NetFix▶onnx(fp32)▶trt(fp32)	13.58ms	0.8779	79.98qps
U2NetFix▶onnx(fp32)▶trt(tf32)	13.31ms	0.8779	84.20qps
U2NetFix▶onnx(fp32)▶trt(fp16)	6.32ms	0.8779	218.06qps
U2NetFix▶onnx(fp32)▶trt(int8)	5.20ms	0.3375	289.77qps

trt모델의 인퍼런스는 아래와 같이 할 수 있다.

import pycuda.autoinit
import pycuda.driver as cuda
import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
trt_runtime = trt.Runtime(TRT_LOGGER)


class HostDeviceMem(object):
    def __init__(self, host_mem, device_mem):
        self.host = host_mem
        self.device = device_mem


def allocate_buffers(engine):
    inputs = []
    outputs = []
    bindings = []
    stream = cuda.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * 1
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        host_mem = cuda.pagelocked_empty(size, dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        bindings.append(int(device_mem))
        if engine.binding_is_input(binding):
            inputs.append(HostDeviceMem(host_mem, device_mem))
        else:
            outputs.append(HostDeviceMem(host_mem, device_mem))
    return inputs, outputs, bindings, stream


def trt_inference(context, bindings, inputs, outputs, stream, img):
    np.copyto(inputs[0].host, img.ravel())
    [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
    stream.synchronize()
    return [out.host.copy() for out in outputs]


trt_model_path = "u2net(fp32_fp16).trt"
with open(trt_model_path, "rb") as f:
    engine_data = f.read()
trt_engine = trt_runtime.deserialize_cuda_engine(engine_data)
inputs, outputs, bindings, stream = allocate_buffers(trt_engine)
context = trt_engine.create_execution_context()

img = load_image(path)

d1, d2, d3, d4, d5, d6, d7 = trt_inference(context, bindings, inputs, outputs, stream, img)

---

시간측정

아래의 방식으로 torch model inference 시간을 측정하라고 권장하나, 이 방법으로는 onnxrt 모델의 추론 시간을 제대로 측정할 수 없다.

starter, ender = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
starter.record()
result = model(img)
ender.record()
torch.cuda.synchronize()
elapsed_time = starter.elapsed_time(ender)

따라서 아래의 방식으로 추론 시간의 Wall Time을 측정한다.

t_beg = time.time()
result = model(img)
torch.cuda.current_stream().synchronize()
t_end = time.time()
elapsed_time = (t_end - t_beg) * 1000