fix ofstream binary issue for windows #585

This commit is contained in:
wang-xinyu 2021-06-16 07:58:13 +00:00
parent f8c537586a
commit 101e60348e
3 changed files with 10 additions and 9 deletions

View File

@ -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;

View File

@ -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;
}
}
}

View File

@ -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;