Improve something.

Author: Iainmon

.gitignore

@@ -43,4 +43,5 @@ build/
 libtorch/
 .cache/
 build-release/
-libtorch_static/
+libtorch_static/
+examples/vgg/**/*.pt

bridge/include/bridge.h

@@ -31,6 +31,8 @@ typedef struct nil_scalar_tensor_t {
 
 float* unsafe(const float* arr);
 bridge_tensor_t load_tensor_from_file(const uint8_t* file_path);
+bridge_tensor_t load_tensor_dict_from_file(const uint8_t* file_path,const uint8_t* tensor_key);
+bridge_tensor_t load_run_model(const uint8_t* model_path, bridge_tensor_t input);
 
 
 int baz(void);

bridge/lib/bridge.cpp

@@ -1,6 +1,8 @@
 #include <bridge.h>
 
 #include <torch/torch.h>
+#include <torch/script.h>
+
 // #include <torch/script.h>
 // #include <Aten/ATen.h>
 #include <iostream>
@@ -78,26 +80,62 @@ std::vector<char> get_the_bytes(std::string filename) {
 extern "C" bridge_tensor_t load_tensor_from_file(const uint8_t* file_path) {
     // // Load the tensor from a file
     // torch::Tensor tensor;
-    // // torch::load(tensor,file_path);
+    // torch::load(tensor,file_path);
 
     // std::cout << "Tensor loaded from file: " << tensor.sizes() << std::endl;
 
-    // // Convert the tensor to a bridge_tensor_t
+    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();
+    return torch_to_bridge(t);
+}
+
+extern "C" bridge_tensor_t load_tensor_dict_from_file(const uint8_t* file_path,const uint8_t* tensor_key) {
+    
+
+    std::cout << "Loading tensor from file: " << file_path << std::endl;
+    std::cout << "Tensor key: " << tensor_key << std::endl;
+
+    std::cout.flush();
 
     std::string fp(reinterpret_cast<const char*>(file_path));
-    std::cout << "File path: " << fp << std::endl;
+    std::string tk(reinterpret_cast<const char*>(tensor_key));
 
-    std::vector<char> f = get_the_bytes(fp);
-    std::cout << "File size: " << f.size() << std::endl;
+    torch::jit::script::Module container = torch::jit::load(fp);
+    torch::Tensor tensor = container.attr(tk).toTensor();
 
-    torch::IValue x = torch::pickle_load(f);
-    // std::cout << "IValue loaded from file: " << x << std::endl;
+    return torch_to_bridge(tensor);
 
-    torch::Tensor t = x.toTensor();
-    std::cout << "Tensor loaded from IValue: " << t.sizes() << std::endl;
-    std::cout << "Tensor sum: " << t.sum() << std::endl;
+}
 
-    return torch_to_bridge(t);
+extern "C" bridge_tensor_t load_run_model(const uint8_t* model_path, bridge_tensor_t input) {
+    auto t_input = bridge_to_torch(input);
+    std::string mp(reinterpret_cast<const char*>(model_path));
+
+    std::cout << "Loading model from path: " << mp << std::endl;
+    std::cout.flush();
+
+
+    torch::jit::Module module;
+    try
+    {
+        // Deserialize the ScriptModule from a file using torch::jit::load().
+        module = torch::jit::load(mp);
+    }
+    catch (const c10::Error& e)
+    {
+        std::cerr << "error loading the model\n" << e.msg();
+        std::system("pause");
+    }
+
+    std::vector<torch::jit::IValue> inputs;
+    inputs.push_back(t_input);
+
+    auto output = module.forward(inputs).toTensor();
+
+    std::cout << "Model output: " << output.sizes() << std::endl;
+    return torch_to_bridge(output);
 }
 
 
@@ -108,6 +146,7 @@ extern "C" bridge_tensor_t load_tensor_from_file(const uint8_t* file_path) {
 
 
 
+
 extern "C" bridge_tensor_t increment3(bridge_tensor_t arr) {
     auto t = bridge_to_torch(arr);
     // Increment the tensor

examples/vgg/dump_weights.py

@@ -3,20 +3,71 @@
 import sys
 import os
 
+import vgg
 
-from vgg import vgg16
 
-model = vgg16(pretrained=True)
+# from vgg import vgg16
 
-model = model.to(torch.float16)
+# model = vgg16(pretrained=True)
 
-# Add the scripts directory to the sys.path
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'scripts')))
+# model = model.to(torch.float16)
 
-import chai
+# # Add the scripts directory to the sys.path
+# sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'scripts')))
 
-os.makedirs('models/vgg16', exist_ok=True)
-model.chai_dump('models/vgg16','vgg16', with_json=False, verbose=True)
+# import chai
 
-# print(model.state_dict().keys())
-# print([(n,w.dtype) for (n,w) in model.state_dict().items()])
+# os.makedirs('models/vgg16', exist_ok=True)
+# model.chai_dump('models/vgg16','vgg16', with_json=False, verbose=True)
+
+# # print(model.state_dict().keys())
+# # print([(n,w.dtype) for (n,w) in model.state_dict().items()])
+
+
+for vggxx in vgg.model_urls.keys():
+    print(vggxx)
+    model = vgg.__dict__[vggxx](pretrained=True)
+    model = model.to(torch.float32)
+    model.eval()
+
+    # # Add the scripts directory to the sys.path
+    # sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'scripts')))
+
+    # import chai
+
+    # os.makedirs('models/vgg16', exist_ok=True)
+    # model.chai_dump('models/vgg16',vggxx, with_json=False, verbose=True)
+
+    # print(f'models/traced_{vggxx}.pt')
+    # faux_input = torch.randn(3,720,1280)
+    # traced_model = torch.jit.trace(model, faux_input)
+    # traced_model.save(f'models/traced_{vggxx}.pt')
+
+    print(f'models/{vggxx}.pt')
+    sd = model.state_dict()
+    torch.save(sd, f'models/{vggxx}.pt')
+
+    # print(f'models/traced_{vggxx}.pt')
+    # faux_input = torch.randn(3,720,1280).to(torch.float16)
+    # traced_model = torch.jit.trace(model, faux_input)
+    # torch.save(traced_model,f'models/traced_{vggxx}.pt')
+
+    print(f'models/script_{vggxx}.pt')
+    script_model = torch.jit.script(model)
+    script_model.save(f'models/script_{vggxx}.pt')
+
+    test = torch.rand(1,3,720,1280).to(torch.float32)
+    # print(model(test).shape)
+
+    model.eval()
+    print(f'models/trace_{vggxx}.pt')
+    faux_input = torch.rand(1,3,720,1280).to(torch.float32)
+    print(model(faux_input).shape)
+    trace_model = torch.jit.trace(model,faux_input)
+    trace_model.save(f'models/trace_{vggxx}.pt')
+    # torch.save(trace_model,f'models/trace_{vggxx}.pt')
+
+
+    # print(model.state_dict().keys())
+    # print([(n,w.dtype) for (n,w) in model.state_dict().items()])
+    

examples/vgg/mktensor.ipynb

@@ -78,11 +78,37 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 24,
    "id": "1a17ec15",
    "metadata": {},
    "outputs": [],
-   "source": []
+   "source": [
+    "m = {'a': x, 'b': x + 1}\n",
+    "torch.save(m, 'my_tensor_dict.pt')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "841d4dfe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class MyModule(torch.nn.Module):\n",
+    "    def __init__(self):\n",
+    "        super(MyModule, self).__init__()\n",
+    "        self.conv1 = torch.nn.Conv2d(in_channels=3,\n",
+    "                                     out_channels=16,\n",
+    "                                     kernel_size=3,\n",
+    "                                     stride=2)\n",
+    "        self.conv2 = torch.nn.Conv2d(in_channels=16,\n",
+    "                                     out_channels=32,\n",
+    "                                     kernel_size=3,\n",
+    "                                     stride=2)\n",
+    "\n",
+    "    def forward(self, x):\n",
+    "        return self.linear(x)"
+   ]
   }
  ],
  "metadata": {

examples/vgg/models/vgg16/classifier.0.bias.meta.json

@@ -5,6 +5,6 @@
   ],
   "rank": 1,
   "size": 4096,
-  "dtype": "float16",
-  "element_bits": 16
+  "dtype": "float32",
+  "element_bits": 32
 }

examples/vgg/models/vgg16/classifier.0.weight.meta.json

@@ -6,6 +6,6 @@
   ],
   "rank": 2,
   "size": 102760448,
-  "dtype": "float16",
-  "element_bits": 16
+  "dtype": "float32",
+  "element_bits": 32
 }

examples/vgg/models/vgg16/classifier.3.bias.meta.json

@@ -5,6 +5,6 @@
   ],
   "rank": 1,
   "size": 4096,
-  "dtype": "float16",
-  "element_bits": 16
+  "dtype": "float32",
+  "element_bits": 32
 }

examples/vgg/models/vgg16/classifier.3.weight.meta.json

@@ -6,6 +6,6 @@
   ],
   "rank": 2,
   "size": 16777216,
-  "dtype": "float16",
-  "element_bits": 16
+  "dtype": "float32",
+  "element_bits": 32
 }

examples/vgg/models/vgg16/classifier.6.bias.meta.json

@@ -5,6 +5,6 @@
   ],
   "rank": 1,
   "size": 1000,
-  "dtype": "float16",
-  "element_bits": 16
+  "dtype": "float32",
+  "element_bits": 32
 }