Fixed torch tensor loading.

Author: Iainmon

CMakeLists.txt

@@ -216,6 +216,7 @@ add_dependencies(TorchBridge ChAI)
 add_dependencies(TorchBridge bridge_objs)
 target_link_options(TorchBridge
     PRIVATE
+        -M ${PROJECT_ROOT_DIR}/lib
         ${BRIDGE_DIR}/include/bridge.h
         ${BRIDGE_OBJECT_FILES}
         -L ${LIBTORCH_DIR}/lib
@@ -225,6 +226,7 @@ target_link_options(TorchBridge
         "-ltorch_global_deps"
         ${LIBTORCH_LIBS_LINKER_ARGS}
         --ldflags "-Wl,-rpath,${LIBTORCH_DIR}/lib"
+        ${CHAI_LINKER_ARGS}
 )
 
 

bridge/lib/bridge.cpp

@@ -94,7 +94,8 @@ void store_tensor(at::Tensor &input, float32_t* dest) {
     std::memcpy(dest,data,bytes_size);
 }
 
-bridge_tensor_t torch_to_bridge(at::Tensor &tensor) {
+bridge_tensor_t torch_to_bridge(at::Tensor tensor_) {
+    at::Tensor tensor = tensor_.contiguous().to(torch::kCPU,torch::kFloat32,false,false);
     bridge_tensor_t result;
     result.created_by_c = true;
     result.was_freed    = false;
@@ -170,7 +171,9 @@ extern "C" bridge_tensor_t load_tensor_from_file(const uint8_t* file_path) {
     std::string fp(reinterpret_cast<const char*>(file_path));
     std::vector<char> f = get_the_bytes(fp);
     torch::IValue x = torch::pickle_load(f);
-    torch::Tensor t = x.toTensor();
+    at::Tensor t = x.toTensor();
+    std::cout << "Tensor loaded from file: " << t.sizes() << std::endl;
+    std::cout.flush();
     return torch_to_bridge(t);
 }
 

examples/CMakeLists.txt

@@ -3,6 +3,3 @@
 add_subdirectory(my_example)
 
 add_subdirectory(torch_model_loading)
-
-
-add_subdirectory(split_loop)

examples/split_loop/CMakeLists.txt

@@ -9,14 +9,21 @@ add_executable(TorchLoad
 )
 add_dependencies(TorchLoad bridge)
 add_dependencies(TorchLoad ChAI)
+add_dependencies(TorchLoad bridge_objs)
 target_link_options(TorchLoad
     PRIVATE
-    ${CHAI_LINKER_ARGS}
+      --main-module torch_load.chpl
+      -M ${PROJECT_ROOT_DIR}/lib
+      ${CHAI_LINKER_ARGS}
 )
 
 add_custom_command(TARGET TorchLoad POST_BUILD
   COMMAND ${CMAKE_COMMAND} -E copy_directory
           "${CMAKE_CURRENT_SOURCE_DIR}/models"
-          "$<TARGET_FILE_DIR:TorchLoad>/models"
+          "${CMAKE_BINARY_DIR}/TorchLoad_models"
   COMMENT "Copying model folder"
+)
+
+set_target_properties(TorchLoad PROPERTIES
+    RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}
 )

examples/split_loop/split_loop.chpl

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "id": "873dd3b8",
    "metadata": {},
    "outputs": [],
@@ -12,7 +12,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "a07c23ff",
    "metadata": {},
    "outputs": [],
@@ -36,7 +36,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "131adc46",
    "metadata": {},
    "outputs": [
@@ -57,39 +57,49 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 11,
    "id": "d4aed442",
    "metadata": {},
    "outputs": [],
    "source": [
-    "x = torch.arange(0,num_elt)\n",
-    "x = x.reshape(f1,f2).to(torch.float32)"
+    "x = torch.arange(0,(500 ** 2)*3)\n",
+    "x = x.reshape(3,500,500).to(torch.float32)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
-   "id": "481d4709",
+   "execution_count": 12,
+   "id": "c45c09fc",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([3, 500, 500])"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
-    "torch.save(x, 'models/my_tensor.pt')"
+    "x.shape"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "id": "1a17ec15",
+   "execution_count": 13,
+   "id": "481d4709",
    "metadata": {},
    "outputs": [],
    "source": [
-    "m = {'a': x, 'b': x + 1}\n",
-    "torch.save(m, 'models/my_tensor_dict.pt')"
+    "torch.save(x, 'models/input.pt')"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 9,
    "id": "841d4dfe",
    "metadata": {},
    "outputs": [],
@@ -112,7 +122,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 10,
    "id": "ca34dafc",
    "metadata": {},
    "outputs": [
@@ -122,7 +132,7 @@
        "True"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }