[fixes] Release memory.

MatzeB · Dec 7, 2018 · 3a582e3 · 3a582e3
1 parent f15c7f4
commit 3a582e3
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 17 deletions.
diff --git a/lib/Importer/Caffe2ModelLoader.cpp b/lib/Importer/Caffe2ModelLoader.cpp
@@ -846,8 +846,7 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
      */
 
     for (auto &o : op.output()) {
-      auto *T = new Tensor();
-      tensors_[o].reset(T);
+      std::unique_ptr<Tensor> T(new Tensor());
 
       auto dim = getShape(dict["shape"]);
 
@@ -884,6 +883,7 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
       RETURN_ERR_IF_NOT(i == T->size(),
                         "The number of serialized values does not "
                         "match the size of the tensor.");
+      tensors_[o] = std::move(T);
     }
 
     return llvm::Error::success();
@@ -917,11 +917,10 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
      }
      */
     for (auto &o : op.output()) {
-      auto *T = new Tensor();
+      std::unique_ptr<Tensor> T(new Tensor());
       if (tensors_.count(o)) {
         continue;
       }
-      tensors_[o].reset(T);
 
       auto dim = getShape(dict["shape"]);
 
@@ -956,6 +955,8 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
       RETURN_ERR_IF_NOT(i == T->size(),
                         "The number of serialized values does not "
                         "match the size of the tensor.");
+
+      tensors_[o] = std::move(T);
     }
 
     return llvm::Error::success();
@@ -980,8 +981,7 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
       return llvm::Error::success();
     }
 
-    auto *T = new Tensor();
-    tensors_[name].reset(T);
+    std::unique_ptr<Tensor> T(new Tensor());
 
     // The shape is set either the shape argument, or from another input
     // tensor. Shape takes priority over input.
@@ -1030,6 +1030,7 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
       RETURN_ERR("Unsupported datatype for ConstantFill.");
     }
 
+    tensors_[name] = std::move(T);
     return llvm::Error::success();
   }
 
@@ -1053,8 +1054,7 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
      }
     */
     const auto &name = op.output(0);
-    auto *T = new Tensor();
-    tensors_[name].reset(T);
+    std::unique_ptr<Tensor> T(new Tensor());
     auto dim = getShape(dict["shape"]);
     T->reset(ElemKind::FloatTy, dim);
     auto TH = T->getHandle<>();
@@ -1071,6 +1071,8 @@ llvm::Error Caffe2ModelLoader::loadWeight(const caffe2::OperatorDef &op) {
     for (size_t i = 0, e = T->size(); i != e; i++) {
       TH.raw(i) = G_.getParent()->getPRNG().nextRandReal(tensorMin, tensorMax);
     }
+
+    tensors_[name] = std::move(T);
     return llvm::Error::success();
   }
 

diff --git a/lib/Importer/ONNXIFIModelLoader.cpp b/lib/Importer/ONNXIFIModelLoader.cpp
@@ -116,9 +116,9 @@ static llvm::Error loadWeight(const onnxTensorDescriptorV1 &in, Tensor *T) {
 llvm::Error ONNXIFIModelLoader::loadWeights(
     uint32_t weightsCount, const onnxTensorDescriptorV1 *weightDescriptors) {
   for (uint32_t i = 0; i < weightsCount; ++i) {
-    auto *T = new Tensor();
-    tensors_[weightDescriptors[i].name].reset(T);
-    RETURN_IF_ERR(loadWeight(weightDescriptors[i], T));
+    std::unique_ptr<Tensor> T(new Tensor());
+    RETURN_IF_ERR(loadWeight(weightDescriptors[i], T.get()));
+    tensors_[weightDescriptors[i].name] = std::move(T);
   }
 
   return llvm::Error::success();

diff --git a/lib/Importer/ONNXModelLoader.cpp b/lib/Importer/ONNXModelLoader.cpp
@@ -258,9 +258,9 @@ llvm::Error ONNXModelLoader::loadConstant(const ONNX_NAMESPACE::NodeProto &op,
                         ONNX_NAMESPACE::AttributeProto::TENSOR,
                     "Only Tensor type constants are supported.");
 
-  auto *T = new Tensor();
-  tensors_[name].reset(T);
-  RETURN_IF_ERR(loadTensor(dict.at("value")->t(), T));
+  std::unique_ptr<Tensor> T(new Tensor());
+  RETURN_IF_ERR(loadTensor(dict.at("value")->t(), T.get()));
+  tensors_[name] = std::move(T);
 
   return llvm::Error::success();
 }
@@ -828,9 +828,9 @@ llvm::Error ONNXModelLoader::loadOperator(const ONNX_NAMESPACE::NodeProto &op) {
 llvm::Error ONNXModelLoader::loadInitializers(ONNX_NAMESPACE::GraphProto &net) {
   // Load the network initializaers:
   for (const auto &in : net.initializer()) {
-    Tensor *T = new Tensor();
-    tensors_[in.name()].reset(T);
-    RETURN_IF_ERR(loadTensor(in, T));
+    std::unique_ptr<Tensor> T(new Tensor());
+    RETURN_IF_ERR(loadTensor(in, T.get()));
+    tensors_[in.name()] = std::move(T);
   }
   return llvm::Error::success();
 }