简体   繁体   中英

How to concatenate LibTorch tensors created with a multi-thread process std::thread in C++?

A multi-thread process in C++ returns tensors and I want to concatenate them into one tensor in order.

In C++ I have a single function which returns a single 1x8 tensor.
I call this function multiple times simultaneously with std::thread and I want to concatenate the tensors it returns into one large tensor. For example, I call it 12 times I expect to have a 12x8 tensor when it's finished.

I need for them to be concatenated in order, that is, the tensor called with 0 should always go in the 0th position followed by the 1st in the 1st position, and so on.

I'm aware I could just have the function return a 12x8 tensor, but I need to solve the problem of how to grab the tensors produced during the multithreading process.

In my attempt below I try to concatenate the tensors into the all_episode_steps tensor but this returns an error.
If you comment out the all_episode_steps line and put std::cout << one; in the get_tensors function above the return statement you see it appears to be using multithreading to create the tensors with no problems.

#include <torch/torch.h>

torch::Tensor get_tensors(int id) {
    torch::Tensor one = torch::rand({8});
    return one.unsqueeze(0);
}

torch::Tensor all_episode_steps;
int main() {

    std::thread ths[100];

    for (int id=0; id<12; id++) {

        ths[id] = std::thread(get_tensors, id);

        all_episode_steps = torch::cat({ths[id], all_episode_steps});
    }
    for (int id=0; id<12; id++) {
        ths[id].join();
    }

}

If you want to build this yourself you can install LibTorch here .
Below is the CMakeLists.txt file for the code above.

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example-app)

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 14)

# The following code block is suggested to be used on Windows.
# According to https://github.com/pytorch/pytorch/issues/25457,
# the DLLs need to be copied to avoid memory errors.
if (MSVC)
  file(GLOB TORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll")
  add_custom_command(TARGET example-app
                     POST_BUILD
                     COMMAND ${CMAKE_COMMAND} -E copy_if_different
                     ${TORCH_DLLS}
                     $<TARGET_FILE_DIR:example-app>)
endif (MSVC)

Threads can't return tensors, but they can modify tensors via pointers. Try this (untested, may need a bit of tweaking):

void get_tensors(torch::Tensor* out) {
    torch::Tensor one = torch::rand({8});
    *out = one.unsqueeze(0);
}

int main() {
    std::thread ths[12];
    std::vector<torch::Tensor> results(12);

    for (int id=0; id<12; id++) {
        ths[id] = std::thread(get_tensors, &results[id]);
    }

    for (int id=0; id<12; id++) {
        ths[id].join();
    }

    auto result2d = torch::cat(results);
}

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM