05. TensorRT VGG

VGG 모델을 TensorRT로 포팅을 진행해보겠습니다.

 

Weights 추출하기

먼저 TensorRT 에서 모델을 사용하기 위해서는 Weights 파일을 추출해야 합니다.

 

import torch
import struct
from torchsummary import summary


def main():
    print('cuda device count: ', torch.cuda.device_count())
    net = torch.load('vgg.pth')
    net = net.to('cuda:0')
    net = net.eval()
    print('model: ', net)
    # print('state dict: ', net.state_dict().keys())
    tmp = torch.ones(1, 3, 224, 224).to('cuda:0')
    print('input: ', tmp)
    out = net(tmp)

    print('output:', out)

    summary(net, (3, 224, 224))
    # return
    f = open("vgg.wts", 'w')
    f.write("{}\n".format(len(net.state_dict().keys())))
    for k, v in net.state_dict().items():
        print('key: ', k)
        print('value: ', v.shape)
        vr = v.reshape(-1).cpu().numpy()
        f.write("{} {}".format(k, len(vr)))
        for vv in vr:
            f.write(" ")
            f.write(struct.pack(">f", float(vv)).hex())
        f.write("\n")


if __name__ == '__main__':
    main()

※ 출처 : https://github.com/wang-xinyu/pytorchx/blob/master/vgg/inference.py

 

먼저 Pytorch에서 제공하는 Pretrained VGG를 불러옵니다.

 

위 코드에서는 .pth 로 불러왔지만 실제로 테스트 할때는 torchvision.models 에서 Pretrained VGG11을 가져왔습니다.

 

TensorRT에서 사용하기 위해서 key, value로 나누어 .wts 형식으로 저장해줍니다.

 

이중 하나를 print 해보면 아래와 같습니다.

• key:  features.0.weight
• value:  torch.Size([64, 3, 3, 3])

또한 Value 값을 저장해 줄 때 빅라디안 방식의 float 형식, 16진수로 값을 하나씩 띄어쓰기를 기준으로 저장해줍니다.

 

 

저장되어 있는 형태를 간단히 보자면 위와 같습니다.

 

※ 22개의 key, value 정보

key:  features.0.weight   value:  torch.Size([64, 3, 3, 3])
key:  features.0.bias   value:  torch.Size([64])
key:  features.3.weight   value:  torch.Size([128, 64, 3, 3])
key:  features.3.bias   value:  torch.Size([128])
key:  features.6.weight   value:  torch.Size([256, 128, 3, 3])
key:  features.6.bias   value:  torch.Size([256])
key:  features.8.weight   value:  torch.Size([256, 256, 3, 3])
key:  features.8.bias   value:  torch.Size([256])
key:  features.11.weight   value:  torch.Size([512, 256, 3, 3])
key:  features.11.bias   value:  torch.Size([512])
key:  features.13.weight   value:  torch.Size([512, 512, 3, 3])
key:  features.13.bias   value:  torch.Size([512])
key:  features.16.weight   value:  torch.Size([512, 512, 3, 3])
key:  features.16.bias   value:  torch.Size([512])
key:  features.18.weight   value:  torch.Size([512, 512, 3, 3])
key:  features.18.bias   value:  torch.Size([512])
key:  classifier.0.weight   value:  torch.Size([4096, 25088])
key:  classifier.0.bias   value:  torch.Size([4096])
key:  classifier.3.weight   value:  torch.Size([4096, 4096])
key:  classifier.3.bias   value:  torch.Size([4096])
key:  classifier.6.weight   value:  torch.Size([1000, 4096])
key:  classifier.6.bias   value:  torch.Size([1000])

 

---

TensorRT 사용하기

 

이제 TensorRT를 사용해보겠습니다.

 

먼저 모델에 관련된 변수들을 선언해줍니다.

 

// 1. 변수 선언
input, output data shape
static const int INPUT_H = 224;
static const int INPUT_W = 224;
static const int INPUT_C = 3;
static const int OUTPUT_SIZE = 1000;
unsigned int batch_size = 1;		// 사용할 배치 사이즈 값

bool serialize = false;				// Engine 생성 유무, true : 엔진 생성
char engineFileName[] = "vgg11";	// Engine 이름

const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";

char engine_file_path[256];
sprintf(engine_file_path, "../Engine/%s.engine", engineFileName);	// Engine 저장 경로

• Input 이미지의 H, W, C
• Output Size
• Batch Size
• Serialize : Engine 강제 생성 유무
• Engine Name
• Input Name
• Output Name

 

// Builds an engine from a network definition.
IBuilder* builder = createInferBuilder(gLogger);

// Holds properties for configuring a builder to produce an engine.
IBuilderConfig* config = builder->createBuilderConfig();

// Create Engine
createEngine(batch_size, builder, config, DataType::kFLOAT, engine_file_path);

// Destroy builder, config
builder->destroy();
config->destroy();

• Builder, Config : 엔진 생성을 위한 변수

엔진 생성을 위해 createEngine 함수를 만들어 줍니다.

 

// Create Engine

std::cout << "------------ Model Build Start ------------" << std::endl;
// A network definition for input to the builder.
INetworkDefinition* network = builder->createNetworkV2(0U);

// Load .wts File
// Weights : [DataType type, const void* values, int64_t count]
std::map<std::string, Weights> weightMap = loadWeights("../Weights/vgg11.wts");

// Create Input Data
ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{ INPUT_H, INPUT_W, INPUT_C });
assert(data);

// Create Model Layers
IConvolutionLayer* conv1 = network->addConvolutionNd(*data, 64, DimsHW{ 3, 3 }, weightMap["features.0.weight"], weightMap["features.0.bias"]);

...

fc1->getOutput(0)->setName(OUTPUT_BLOB_NAME);
std::cout << "set name out" << std::endl;
network->markOutput(*fc1->getOutput(0));

// Build engine
builder->setMaxBatchSize(batch_size);   // 사용할 Batch Size 설정
config->setMaxWorkspaceSize(1 << 20);   // 사용할 메모리 크기 설정

std::cout << "Building engine, please wait for a while..." << std::endl;
IHostMemory* engine = builder->buildSerializedNetwork(*network, *config);
std::cout << "==== model build done ====" << std::endl << std::endl;

std::cout << "==== model selialize start ====" << std::endl << std::endl;
std::ofstream p(engineFileName, std::ios::binary);
if (!p) {
    std::cerr << "could not open plan output file" << std::endl << std::endl;
}
p.write(reinterpret_cast<const char*>(engine->data()), engine->size());
std::cout << "==== model selialize done ====" << std::endl << std::endl;

engine->destroy();
network->destroy();
p.close();
// Release host memory
for (auto& mem : weightMap)
{
    free((void*)(mem.second.values));
}

모델 Layer는 너무 길어서 생략했습니다.

 

이제 만들거나 불러온 Engine File이 올바른지 확인을 해봅니다.

 

// 4. Engine file 로드 
char* trtModelStream{ nullptr };// 저장된 스트림을 저장할 변수
size_t size{ 0 };
std::cout << "------------ Engine file load ------------" << std::endl << std::endl;
std::ifstream file(engine_file_path, std::ios::binary);
if (file.good()) {							// good() : 입출력 가능 여부 확인
    file.seekg(0, file.end);				// seekg() : 포인터 이동
    size = file.tellg();					// tellg() : 크기
    file.seekg(0, file.beg);
    trtModelStream = new char[size];
    file.read(trtModelStream, size);
    file.close();
}
else {
    std::cout << "[ERROR] Engine file load error" << std::endl;
}

Engine File을 읽어와 trModelStream 변수에 저장시켜 주고

 

// 5. Engine File 로드 후 Engine 생성
std::cout << "------------ Engine file deserialize ------------" << std::endl << std::endl;
IRuntime* runtime = createInferRuntime(gLogger);
ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size);
IExecutionContext* context = engine->createExecutionContext();
delete[] trtModelStream;

void* buffers[2];
const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME);
const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME);

// GPU에서 입력과 출력으로 사용할 메모리 공간할당
CHECK(cudaMalloc(&buffers[inputIndex], batch_size * INPUT_C * INPUT_H * INPUT_W * sizeof(uint8_t)));
CHECK(cudaMalloc(&buffers[outputIndex], batch_size * OUTPUT_SIZE * sizeof(float)));

Engine을 생성하게 됩니다.

 

이제 데이터를 넣어 연산을 진행해봅니다.

 

// 7. Inference
// CUDA 스트림 생성
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));

for (int i = 0; i < iter_count; i++) {
    auto start = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count();

    CHECK(cudaMemcpyAsync(buffers[inputIndex], input.data(), batch_size * INPUT_C * INPUT_H * INPUT_W * sizeof(uint8_t), cudaMemcpyHostToDevice, stream));
    context->enqueue(batch_size, buffers, stream, nullptr);
    CHECK(cudaMemcpyAsync(outputs.data(), buffers[outputIndex], batch_size * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream));
    cudaStreamSynchronize(stream);

    auto dur = std::chrono::duration_cast<std::chrono::milliseconds>(std::chrono::system_clock::now().time_since_epoch()).count() - start;
    dur_time += dur;
    std::cout << dur << " milliseconds" << std::endl;
}

먼저 GPU를 사용하기 위해 CUDA Stream을 생성해줍니다.

 

CUDA 관련 내용은 다루지 않겠습니다.

 

이제 결과를 출력해주면 끝입니다!

 

// 8. 결과 출력
int max_index = max_element(outputs.begin(), outputs.end()) - outputs.begin();

// Release stream and buffers ...
cudaStreamDestroy(stream);
CHECK(cudaFree(buffers[inputIndex]));
CHECK(cudaFree(buffers[outputIndex]));
context->destroy();
engine->destroy();
runtime->destroy();

 

 

 

 

 

 

 

 

 

※ 참고 자료 : 

https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Core/Logger.html

 

https://github.com/yester31/TensorRT_EX

↑ 굉장히 잘나와있다 추천!