load op optimize

paddle/fluid/operators/load_op.cc

@@ -1,4 +1,5 @@
-
+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
 /* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -35,6 +36,8 @@ class LoadOp : public framework::OperatorBase {
     auto *dev_ctx = platform::DeviceContextPool::Instance().Get(place);
     platform::RecordEvent record_event(Type(), dev_ctx);
 
+    // FIXME(yuyang18): We save variable to local file now, but we should change
+    // it to save an output stream.
     auto filename = Attr<std::string>("file_path");
     std::ifstream fin(filename);
     PADDLE_ENFORCE(static_cast<bool>(fin), "Cannot open file %s for load op",
@@ -46,31 +49,23 @@ class LoadOp : public framework::OperatorBase {
                    out_var_name);
 
     if (out_var->IsType<framework::LoDTensor>()) {
-      LoadLodTensor(filename, place, out_var);
+      LoadLodTensor(fin, place, out_var);
     } else if (out_var->IsType<framework::SelectedRows>()) {
-      LoadSelectedRows(filename, scope, place, out_var);
+      LoadSelectedRows(fin, place, out_var);
     } else {
       PADDLE_ENFORCE(
           false,
           "Load only support LoDTensor and SelectedRows, %s has wrong type",
           out_var_name);
     }
-    }
+  }
 
-  void LoadLodTensor(const std::string &filename, const platform::Place &place,
+  void LoadLodTensor(std::istream &fin, const platform::Place &place,
                      framework::Variable *var) const {
     // get device context from pool
     platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
     auto &dev_ctx = *pool.Get(place);
-
-    // FIXME(yuyang18): We save variable to local file now, but we should change
-    // it to save an output stream.
-    std::ifstream fin(filename);
-    PADDLE_ENFORCE(static_cast<bool>(fin), "Cannot open %s to read",
-                   filename);
-
-      auto *tensor = var->GetMutable<framework::LoDTensor>();
-
+    auto *tensor = var->GetMutable<framework::LoDTensor>();
     DeserializeFromStream(fin, tensor, dev_ctx);
 
     auto load_as_fp16 = Attr<bool>("load_as_fp16");
@@ -92,25 +87,15 @@ class LoadOp : public framework::OperatorBase {
       tensor = var->GetMutable<framework::LoDTensor>();
       tensor->set_lod(fp16_tensor.lod());
       tensor->ShareDataWith(fp16_tensor);
-      }
+    }
   }
 
-  void LoadSelectedRows(const std::string &filename,
-                        const framework::Scope &scope,
-                        const platform::Place &place,
+  void LoadSelectedRows(std::istream &fin, const platform::Place &place,
                         framework::Variable *var) const {
-
     auto *selectedRows = var->GetMutable<framework::SelectedRows>();
-
     // get device context from pool
     platform::DeviceContextPool &pool = platform::DeviceContextPool::Instance();
     auto &dev_ctx = *pool.Get(place);
-
-    // FIXME(yuyang18): We save variable to local file now, but we should change
-    // it to save an output stream.
-    std::ifstream fin(filename);
-    PADDLE_ENFORCE(static_cast<bool>(fin), "Cannot open %s to read",
-                   filename);
     framework::DeserializeFromStream(fin, selectedRows, dev_ctx);
   }
 };