diff --git a/alexnet/alex.cpp b/alexnet/alex.cpp index 17f2a10..f5e649b 100644 --- a/alexnet/alex.cpp +++ b/alexnet/alex.cpp @@ -233,7 +233,7 @@ int main(int argc, char** argv) APIToModel(1, &modelStream); assert(modelStream != nullptr); - std::ofstream p("alexnet.engine"); + std::ofstream p("alexnet.engine", std::ios::binary); if (!p) { std::cerr << "could not open plan output file" << std::endl; diff --git a/inception/inceptionv4/inception_v4.cpp b/inception/inceptionv4/inception_v4.cpp index 80ea80c..0562fde 100644 --- a/inception/inceptionv4/inception_v4.cpp +++ b/inception/inceptionv4/inception_v4.cpp @@ -44,7 +44,7 @@ namespace trtx { builder -> destroy(); // write serialized engine to file - std::ofstream trtFile(mParams.trtEngineFile); + std::ofstream trtFile(mParams.trtEngineFile, std::ios::binary); if(!trtFile){ std::cerr << "Unable to open engine file." << std::endl; return false; @@ -193,26 +193,26 @@ namespace trtx { // Engine requires exactly IEngine::getNbBindings() number of buffers. assert(mEngine -> getNbBindings() == 2); void* buffers[2]; - + // In order to bind the buffers, we need to know the names of the input and output tensors. // Note that indices are guaranteed to be less than IEngine::getNbBindings() const int inputIndex = mEngine->getBindingIndex(mParams.inputTensorName); const int outputIndex = mEngine->getBindingIndex(mParams.outputTensorName); - + // Create GPU buffers on device CUDA_CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * mParams.inputH * mParams.inputW * sizeof(float))); CUDA_CHECK(cudaMalloc(&buffers[outputIndex], batchSize * 1000 * sizeof(float))); - + // Create stream cudaStream_t stream; CUDA_CHECK(cudaStreamCreate(&stream)); - + // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host CUDA_CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * mParams.inputH * mParams.inputW * sizeof(float), cudaMemcpyHostToDevice, stream)); mContext->enqueue(batchSize, buffers, stream, nullptr); CUDA_CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * 1000 * sizeof(float), cudaMemcpyDeviceToHost, stream)); cudaStreamSynchronize(stream); - + // Release stream and buffers cudaStreamDestroy(stream); CUDA_CHECK(cudaFree(buffers[inputIndex])); @@ -232,4 +232,5 @@ namespace trtx { return true; } -} \ No newline at end of file +} + diff --git a/lenet/lenet.cpp b/lenet/lenet.cpp index 6df3c8e..5fba240 100644 --- a/lenet/lenet.cpp +++ b/lenet/lenet.cpp @@ -221,7 +221,7 @@ int main(int argc, char** argv) APIToModel(1, &modelStream); assert(modelStream != nullptr); - std::ofstream p("lenet5.engine"); + std::ofstream p("lenet5.engine", std::ios::binary); if (!p) { std::cerr << "could not open plan output file" << std::endl;