1. 작업 현황
어제 찍은 27개의 테스트 동영상에 대해서 사람이 검출되는 프레임만 추출하는 작업은 모두 끝이 났다. 어제 돌려놓고 퇴근했더니 오늘 출근하니까 다 돌아져 있었다.
이제 이렇게 추출된 이미지를 투기행위 탐지 모델에 돌리는 작업을 해주어야 한다.
2. 문제상황
새롭게 받은 모델은 그냥 일반 pytorch 모델인 관계로 연산 속도가 매우 느렸다. 다음은 기존 pytorch 모델의 연산 속도를 tqdm으로 표현한 것이다.
평균적으로 1초당 약 한 번의 iteration을 수행하고 있었다. 그 말인 즉슨, 1초에 평균적으로 1번정도 투기 여부를 판단한다는 뜻이다. 총 이미지는 약 5만장 이상으로, 1초에 이미지가 1개씩 들어가다가는 인턴 끝날때까지 계속 이 작업만 하고 있을 것 같았다.
3. 해결
Pytorch To TensorRT
따라서 pytorch 모델을 TensorRT Engine으로 변환해주는 작업을 해주었다. 오전에 차장님께서 잠깐 회사에 출근하셨길래, 차장님께 해당 작업 관련해서 문의를 드렸었다. 차장님께서는, 투기행위 탐지 관련 모델은 멀티 모델이어서 TensorRT로 변환이 불가하고, 그냥 CUDA로 연산해야한다고 하셨다. 연산이 느리더라도 어쩔 수 없는 부분이라고 하셨다. 동영상 하나에서 추출된 프레임을 모두 연산하는 데 약 20~30분정도의 시간이 걸리는데, 총 27개의 영상이 있으니 마냥 기다리고만 있을수는 없었다. 따라서 일단 무작정 TensorRT로 변환하는 작업을 해보기로 하였다.
Import Model
우선 사용할 모델을 불러와주고, 평가모드로 GPU에 올려주었다. 그 이후, 해당 모델에서 사용할 가중치를 넣어주었고, 이 때 strict=True 옵션을 줘서 모든 파라미터에 대해서 strict하게 로딩을 해주었다.
model = 모델_불러오는_코드.eval().cuda()
model.load_state_dict(torch.load('가중치_경로'), strict=True)
Make Dummy data
그리고, input 데이터와 동일한 shape의 더미 데이터를 만들어주었다. model.parameters()로 확인하려고 하였으나, 모델의 층이 너무 깊어서 제대로 확인하기가 힘들었다. 따라서 모델 config 파일에 있는 image_size를 확인해주었다.
dummy_data = torch.zeros(1, 3, 256, 256).cuda()
torch2trt.torch2trt
torch2trt의 torch2trt를 사용하여 model을 TensorRT로 변환해주었다. 변환된 TensorRT의 가중치는 저장하여 추후에도 사용할 수 있게 하였다. 또한, 변환시 더 빠른 추론을 위해 float32로 되어 있는 부분을 float16으로 변환하는 설정도 넣어주었다. float32를 float16으로 변환하지 않으면 속도의 차이가 그렇게 크지 않다고 한다.
model_trt = torch2trt.torch2trt(model, [dummy_data], fp16_mode=True, max_workspace_size=1 << 25)
torch.save(model_trt.state_dict(), '저장경로')
전체 Script
전체 코드이다. 다음과 같이 예외처리를 해주어, TensorRT가 경로에 존재하면 받아오고 그렇지 않으면 새롭게 생성하는 형태의 로직이다.
try:
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('저장경로'))
TensorRT_model = model_trt.eval()
except:
model = 모델.eval().cuda()
model.load_state_dict(torch.load('가중치_경로'), strict=True)
dummy_data = torch.zeros(1, 3, 256, 256).cuda()
model_trt = torch2trt.torch2trt(model, [dummy_data], fp16_mode=True, max_workspace_size=1 << 25)
torch.save(model_trt.state_dict(), '저장경로')
TensorRT_model = model_trt.eval()
성능 향상
기존에 비해 약 20~30배정도의 성능 향상을 이루었다. 기존 모델은 1초에 약 한 장의 사진을 연산하였으나, TensorRT는 1초에 약 20~30장의 사진을 연산한다. 시간으로 보면, 동영상당 약 30분씩 걸리던 연산을 1분 내외로 완료하고 있는 것이다.
모델 삽입
다음과 같은 형태로 TensorRT를 삽입해주었다. main에서 미리 TensorRT를 가져오고, 모든 이미지에 대해서 for문이 돌면 가져온 모델에 이미지가 들어가서 output을 return하는 형식이다.
(...)
if __name__ == '__main__':
(...)
model_trt = TRTModule()
model_trt.load_state_dict(torch.load('저장경로'))
TensorRT_model = model_trt.eval()
(...)
def model_run(image):
output = TensorRT_model(image)
(...)
return output
(...)
files = '이미지들이 저장된 경로'
for file in files:
model_run(file)
4. 한계
TensorRT는 기존 모델을 여러 최적화 과정을 거쳐 NVIDIA의 GPU에 효율적으로 싣는 역할을 한다. 따라서, 최적화 과정을 거쳐 만들어낸 엔진을 binary 형태로 저장하고 사용할때 마다 엔진을 deserialize하여 사용하여야 하는데, 이 과정이 잘 안된다. 따라서 현재는 그냥 TensorRT를 CUDA에서? 사용하고 있다. 아직 이 부분에 대한 개념이 잘 안잡힌 상태여서 내가 TensorRT의 어떤 부분을 어떻게 사용하고 있는지도 잘 모르겠다.
대략적인 방향
우선, TensorRT는 Optimizer와 Runtime Engine으로 구성되어 있다. 완성된 TensorRT에서 Engine 모듈을 분해하여 Stream으로 묶은 뒤, 각 텐서들을 GPU로 할당하고자 하는데 잘 안된다. engine_data를 이진파일로 받아와서 deserialize하려는데 계속 None값이 반환된다. 계속 자료를 찾아봤는데도 잘 안되어서 내일 차장님께 여쭤봐야겠다.
댓글남기기