Merge branch 'develop' of https://github.com/PaddlePaddle/Paddle into fix_quan_hang

test=develop

cmake/external/boost.cmake

@@ -24,7 +24,7 @@ set(BOOST_PROJECT       "extern_boost")
 # So we use 1.41.0 here.
 set(BOOST_VER           "1.41.0")
 set(BOOST_TAR "boost_1_41_0" CACHE STRING "" FORCE)
-set(BOOST_URL "http://paddlepaddledeps.cdn.bcebos.com/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE)
+set(BOOST_URL "http://paddlepaddledeps.bj.bcebos.com/${BOOST_TAR}.tar.gz" CACHE STRING "" FORCE)
 
 MESSAGE(STATUS "BOOST_TAR: ${BOOST_TAR}, BOOST_URL: ${BOOST_URL}")
 

cmake/external/grpc.cmake

@@ -44,7 +44,7 @@ ExternalProject_Add(
     # 3. keep only zlib, cares, protobuf, boringssl under "third_party",
     #    checkout and clean other dirs under third_party
     # 4. remove .git, and package the directory.
-    URL "http://paddlepaddledeps.cdn.bcebos.com/grpc-v1.10.x.tar.gz"
+    URL "http://paddlepaddledeps.bj.bcebos.com/grpc-v1.10.x.tar.gz"
     URL_MD5  "1f268a2aff6759839dccd256adcc91cf"
     PREFIX          ${GRPC_SOURCES_DIR}
     UPDATE_COMMAND  ""

cmake/external/mklml.cmake

@@ -34,7 +34,7 @@ SET(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_RPATH}" "${MKLML_ROOT}/lib")
 SET(TIME_VERSION "2019.0.1.20181227")
 IF(WIN32)
     SET(MKLML_VER "mklml_win_${TIME_VERSION}" CACHE STRING "" FORCE)
-    SET(MKLML_URL "https://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.zip" CACHE STRING "" FORCE)
+    SET(MKLML_URL "https://paddlepaddledeps.bj.bcebos.com/${MKLML_VER}.zip" CACHE STRING "" FORCE)
     SET(MKLML_LIB                 ${MKLML_LIB_DIR}/mklml.lib)
     SET(MKLML_IOMP_LIB            ${MKLML_LIB_DIR}/libiomp5md.lib)
     SET(MKLML_SHARED_LIB          ${MKLML_LIB_DIR}/mklml.dll)
@@ -43,7 +43,7 @@ ELSE()
     #TODO(intel-huying):
     #  Now enable Erf function in mklml library temporarily, it will be updated as offical version later.
     SET(MKLML_VER "Glibc225_vsErf_mklml_lnx_${TIME_VERSION}" CACHE STRING "" FORCE)
-    SET(MKLML_URL "http://paddlepaddledeps.cdn.bcebos.com/${MKLML_VER}.tgz" CACHE STRING "" FORCE)
+    SET(MKLML_URL "http://paddlepaddledeps.bj.bcebos.com/${MKLML_VER}.tgz" CACHE STRING "" FORCE)
     SET(MKLML_LIB                 ${MKLML_LIB_DIR}/libmklml_intel.so)
     SET(MKLML_IOMP_LIB            ${MKLML_LIB_DIR}/libiomp5.so)
     SET(MKLML_SHARED_LIB          ${MKLML_LIB_DIR}/libmklml_intel.so)

paddle/fluid/imperative/layer.cc

@@ -315,6 +315,9 @@ std::map<std::string, std::vector<VarBase*>> OpBase::ApplyGrad() {
       for (size_t i = 0; i < outputs.size(); ++i) {
         framework::Variable* grad = outputs[i]->var_;
         framework::Variable* orig_grad = origin_outputs[i]->var_;
+        VLOG(3) << "AddTo Called with orig_grad is: "
+                << origin_outputs[i]->name_ << " Grad to be added is "
+                << outputs[i]->name_;
         AddTo(grad, orig_grad, place_);
         delete grad;
       }

paddle/fluid/imperative/tracer.cc

@@ -277,6 +277,7 @@ std::set<std::string> Tracer::Trace(OpBase* op, const VarBasePtrMap& inputs,
           VarBase* var = current_vars_map[var_it->second];
           InitGrad(var, prepared_op.GetDeviceContext());
           grad_out_vars.push_back(var->grads_);
+          VLOG(3) << "grads output var name: " << var->name_;
         }
       }
     }

paddle/fluid/inference/api/demo_ci/run.sh

@@ -27,7 +27,7 @@ if [ -d "$TENSORRT_INCLUDE_DIR" -a -d "$TENSORRT_LIB_DIR" ]; then
 fi
 
 PREFIX=inference-vis-demos%2F
-URL_ROOT=http://paddlemodels.cdn.bcebos.com/${PREFIX}
+URL_ROOT=http://paddlemodels.bj.bcebos.com/${PREFIX}
 
 # download vis_demo data
 function download() {

paddle/fluid/inference/tests/api/CMakeLists.txt

@@ -115,14 +115,14 @@ inference_analysis_test(test_analyzer_transformer SRCS analyzer_transformer_test
 # ocr
 set(OCR_INSTALL_DIR "${INFERENCE_DEMO_INSTALL_DIR}/ocr")
 if (NOT EXISTS ${OCR_INSTALL_DIR})
-    inference_download_and_uncompress(${OCR_INSTALL_DIR} "http://paddlemodels.cdn.bcebos.com/" "inference-vis-demos%2Focr.tar.gz")
+    inference_download_and_uncompress(${OCR_INSTALL_DIR} "http://paddlemodels.bj.bcebos.com/" "inference-vis-demos%2Focr.tar.gz")
 endif()
 inference_analysis_api_test_with_refer_result(test_analyzer_ocr ${OCR_INSTALL_DIR} analyzer_vis_tester.cc SERIAL)
 
 # mobilenet with transpose op
 set(MOBILENET_INSTALL_DIR "${INFERENCE_DEMO_INSTALL_DIR}/mobilenet")
 if (NOT EXISTS ${MOBILENET_INSTALL_DIR})
-    inference_download_and_uncompress(${MOBILENET_INSTALL_DIR} "http://paddlemodels.cdn.bcebos.com/" "inference-vis-demos%2Fmobilenet.tar.gz")
+    inference_download_and_uncompress(${MOBILENET_INSTALL_DIR} "http://paddlemodels.bj.bcebos.com/" "inference-vis-demos%2Fmobilenet.tar.gz")
 endif()
 inference_analysis_api_test_with_refer_result(test_analyzer_mobilenet_transpose ${MOBILENET_INSTALL_DIR} analyzer_vis_tester.cc SERIAL)
 

paddle/fluid/inference/tests/test.cmake

 include(ExternalProject)
-set(INFERENCE_URL "http://paddle-inference-dist.cdn.bcebos.com" CACHE STRING "inference download url")
+set(INFERENCE_URL "http://paddle-inference-dist.bj.bcebos.com" CACHE STRING "inference download url")
 set(INFERENCE_DEMO_INSTALL_DIR "${THIRD_PARTY_PATH}/inference_demo" CACHE STRING
     "A path setting inference demo download directories.")
 

paddle/fluid/memory/allocation/CMakeLists.txt

@@ -61,4 +61,6 @@ nv_test(allocation_and_eigen_test SRCS allocation_and_eigen_test.cu DEPS allocat
 
 cc_test(retry_allocator_test SRCS retry_allocator_test.cc DEPS retry_allocator best_fit_allocator locked_allocator cpu_allocator)
 
-cc_test(allocator_facade_test SRCS allocator_facade_test.cc DEPS allocator_facade)
+cc_test(allocator_facade_abs_flags_test SRCS allocator_facade_abs_flags_test.cc DEPS allocator_facade)
+
+cc_test(allocator_facade_frac_flags_test SRCS allocator_facade_frac_flags_test.cc DEPS allocator_facade)

paddle/fluid/memory/allocation/allocator_facade_abs_flags_test.cc

+// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/memory/allocation/allocator_facade.h"
+#include <gflags/gflags.h>
+#include <gtest/gtest.h>
+
+#ifdef PADDLE_WITH_CUDA
+DECLARE_double(fraction_of_gpu_memory_to_use);
+DECLARE_double(fraction_of_cuda_pinned_memory_to_use);
+DECLARE_uint64(initial_gpu_memory_in_mb);
+DECLARE_uint64(reallocate_gpu_memory_in_mb);
+DECLARE_int64(gpu_allocator_retry_time);
+#endif
+
+namespace paddle {
+namespace memory {
+namespace allocation {
+
+//! Run allocate test cases for different places
+void AllocateTestCases() {
+  auto &instance = AllocatorFacade::Instance();
+  platform::Place place;
+  size_t size = 1024;
+
+  {
+    place = platform::CPUPlace();
+    size = 1024;
+    auto cpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(cpu_allocation, nullptr);
+    ASSERT_NE(cpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(cpu_allocation->place(), place);
+    ASSERT_EQ(cpu_allocation->size(), size);
+  }
+
+#ifdef PADDLE_WITH_CUDA
+  {
+    place = platform::CUDAPlace(0);
+    size = 1024;
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    // Allocate 2GB gpu memory
+    place = platform::CUDAPlace(0);
+    size = 2 * static_cast<size_t>(1 << 30);
+    auto gpu_allocation = instance.Alloc(place, size);
+    ASSERT_NE(gpu_allocation, nullptr);
+    ASSERT_NE(gpu_allocation->ptr(), nullptr);
+    ASSERT_EQ(gpu_allocation->place(), place);
+    ASSERT_GE(gpu_allocation->size(), size);
+  }
+
+  {
+    place = platform::CUDAPinnedPlace();
+    size = (1 << 20);
+    auto cuda_pinned_allocation =
+        instance.Alloc(platform::CUDAPinnedPlace(), 1 << 20);
+    ASSERT_NE(cuda_pinned_allocation, nullptr);
+    ASSERT_NE(cuda_pinned_allocation->ptr(), nullptr);
+    ASSERT_EQ(cuda_pinned_allocation->place(), place);
+    ASSERT_GE(cuda_pinned_allocation->size(), size);
+  }
+#endif
+}
+
+TEST(Allocator, SpecifyGpuMemory) {
+#ifdef PADDLE_WITH_CUDA
+  // Set to 0.0 to test FLAGS_initial_gpu_memory_in_mb and
+  // FLAGS_reallocate_gpu_memory_in_mb
+  FLAGS_fraction_of_gpu_memory_to_use = 0.0;
+  // 512 MB
+  FLAGS_initial_gpu_memory_in_mb = 512;
+  // 4 MB
+  FLAGS_reallocate_gpu_memory_in_mb = 4;
+  FLAGS_gpu_allocator_retry_time = 500;
+  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
+#endif
+
+  AllocateTestCases();
+}
+
+}  // namespace allocation
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/allocation/allocator_facade_test.cc → paddle/fluid/memory/allocation/allocator_facade_frac_flags_test.cc

@@ -19,6 +19,8 @@
 #ifdef PADDLE_WITH_CUDA
 DECLARE_double(fraction_of_gpu_memory_to_use);
 DECLARE_double(fraction_of_cuda_pinned_memory_to_use);
+DECLARE_uint64(initial_gpu_memory_in_mb);
+DECLARE_uint64(reallocate_gpu_memory_in_mb);
 DECLARE_int64(gpu_allocator_retry_time);
 #endif
 
@@ -26,13 +28,8 @@ namespace paddle {
 namespace memory {
 namespace allocation {
 
-TEST(allocator, allocator) {
-#ifdef PADDLE_WITH_CUDA
-  FLAGS_fraction_of_gpu_memory_to_use = 0.01;
-  FLAGS_gpu_allocator_retry_time = 500;
-  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
-#endif
-
+//! Run allocate test cases for different places
+void AllocateTestCases() {
   auto &instance = AllocatorFacade::Instance();
   platform::Place place;
   size_t size = 1024;
@@ -82,6 +79,16 @@ TEST(allocator, allocator) {
 #endif
 }
 
+TEST(Allocator, Allocator) {
+#ifdef PADDLE_WITH_CUDA
+  FLAGS_fraction_of_gpu_memory_to_use = 0.01;
+  FLAGS_gpu_allocator_retry_time = 500;
+  FLAGS_fraction_of_cuda_pinned_memory_to_use = 0.5;
+#endif
+
+  AllocateTestCases();
+}
+
 }  // namespace allocation
 }  // namespace memory
 }  // namespace paddle

paddle/fluid/memory/allocation/legacy_allocator.cc

@@ -37,6 +37,8 @@ DEFINE_bool(init_allocated_mem, false,
             "that initializing the allocated memory with a small value "
             "during unit testing.");
 DECLARE_double(fraction_of_gpu_memory_to_use);
+DECLARE_uint64(initial_gpu_memory_in_mb);
+DECLARE_uint64(reallocate_gpu_memory_in_mb);
 DECLARE_bool(benchmark);
 
 namespace paddle {
@@ -153,12 +155,18 @@ BuddyAllocator *GetGPUBuddyAllocator(int gpu_id) {
                                     platform::GpuMinChunkSize(),
                                     platform::GpuMaxChunkSize());
 
-      VLOG(10) << "\n\nNOTE: each GPU device use "
-               << FLAGS_fraction_of_gpu_memory_to_use * 100
-               << "% of GPU memory.\n"
-               << "You can set GFlags environment variable '"
-               << "FLAGS_fraction_of_gpu_memory_to_use"
-               << "' to change the fraction of GPU usage.\n\n";
+      VLOG(10) << "\n\nNOTE:\n"
+               << "You can set GFlags environment variable "
+               << "'FLAGS_fraction_of_gpu_memory_to_use' "
+               << "or 'FLAGS_initial_gpu_memory_in_mb' "
+               << "or 'FLAGS_reallocate_gpu_memory_in_mb' "
+               << "to change the memory size for GPU usage.\n"
+               << "Current 'FLAGS_fraction_of_gpu_memory_to_use' value is "
+               << FLAGS_fraction_of_gpu_memory_to_use
+               << ". Current 'FLAGS_initial_gpu_memory_in_mb' value is "
+               << FLAGS_initial_gpu_memory_in_mb
+               << ". Current 'FLAGS_reallocate_gpu_memory_in_mb' value is "
+               << FLAGS_reallocate_gpu_memory_in_mb << "\n\n";
     }
   });
 

paddle/fluid/memory/detail/CMakeLists.txt

@@ -9,3 +9,5 @@ endif(${WITH_GPU})
 cc_test(system_allocator_test SRCS system_allocator_test.cc DEPS system_allocator)
 
 cc_library(buddy_allocator SRCS buddy_allocator.cc DEPS memory_block system_allocator glog)
+
+cc_test(buddy_allocator_test SRCS buddy_allocator_test.cc DEPS buddy_allocator)

paddle/fluid/memory/detail/buddy_allocator.cc

@@ -13,6 +13,10 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/memory/detail/buddy_allocator.h"
+
+#include <algorithm>
+#include <utility>
+
 #include "glog/logging.h"
 
 DEFINE_bool(free_idle_memory, false,
@@ -36,9 +40,10 @@ BuddyAllocator::~BuddyAllocator() {
               "have actually been freed";
   while (!pool_.empty()) {
     auto block = static_cast<MemoryBlock*>(std::get<2>(*pool_.begin()));
-    VLOG(10) << "Free from block (" << block << ", " << max_chunk_size_ << ")";
+    VLOG(10) << "Free from block (" << block << ", " << block->size(cache_)
+             << ")";
 
-    system_allocator_->Free(block, max_chunk_size_, block->index(cache_));
+    system_allocator_->Free(block, block->size(cache_), block->index(cache_));
     cache_.invalidate(block);
     pool_.erase(pool_.begin());
   }
@@ -71,7 +76,7 @@ void* BuddyAllocator::Alloc(size_t unaligned_size) {
 
   // refill the pool if failure
   if (it == pool_.end()) {
-    it = RefillPool();
+    it = RefillPool(size);
     // if still failure, fail fatally
     if (it == pool_.end()) {
       return nullptr;
@@ -184,19 +189,28 @@ void* BuddyAllocator::SystemAlloc(size_t size) {
   return static_cast<MemoryBlock*>(p)->data();
 }
 
-BuddyAllocator::PoolSet::iterator BuddyAllocator::RefillPool() {
+BuddyAllocator::PoolSet::iterator BuddyAllocator::RefillPool(
+    size_t request_bytes) {
+  size_t allocate_bytes = max_chunk_size_;
+  size_t index = 0;
+
 #ifdef PADDLE_WITH_CUDA
   if (system_allocator_->UseGpu()) {
     if ((total_used_ + total_free_) == 0) {
-      // Compute the maximum allocation size for the first allocation.
-      max_chunk_size_ = platform::GpuMaxChunkSize();
+      // Compute the allocation size for gpu for the first allocation.
+      allocate_bytes = std::max(platform::GpuInitAllocSize(), request_bytes);
+    } else {
+      // Reallocation size
+      if (realloc_size_ == 0) {
+        realloc_size_ = platform::GpuReallocSize();
+      }
+      allocate_bytes = std::max(realloc_size_, request_bytes);
     }
   }
 #endif
 
-  // Allocate a new maximum sized block
-  size_t index = 0;
-  void* p = system_allocator_->Alloc(&index, max_chunk_size_);
+  // Allocate a new block
+  void* p = system_allocator_->Alloc(&index, allocate_bytes);
 
   if (p == nullptr) return pool_.end();
 
@@ -204,7 +218,7 @@ BuddyAllocator::PoolSet::iterator BuddyAllocator::RefillPool() {
            << " from system allocator";
 
   static_cast<MemoryBlock*>(p)->init(&cache_, MemoryBlock::FREE_CHUNK, index,
-                                     max_chunk_size_, nullptr, nullptr);
+                                     allocate_bytes, nullptr, nullptr);
 
   // gpu fallback allocation
   if (system_allocator_->UseGpu() &&
@@ -212,10 +226,10 @@ BuddyAllocator::PoolSet::iterator BuddyAllocator::RefillPool() {
     fallback_alloc_count_++;
   }
 
-  total_free_ += max_chunk_size_;
+  total_free_ += allocate_bytes;
 
   // dump the block into pool
-  return pool_.insert(IndexSizeAddress(index, max_chunk_size_, p)).first;
+  return pool_.insert(IndexSizeAddress(index, allocate_bytes, p)).first;
 }
 
 BuddyAllocator::PoolSet::iterator BuddyAllocator::FindExistChunk(size_t size) {
@@ -286,12 +300,12 @@ void BuddyAllocator::CleanIdleFallBackAlloc() {
 
     VLOG(10) << "Return block " << block << " to fallback allocator.";
 
-    system_allocator_->Free(block, max_chunk_size_, block->index(cache_));
+    system_allocator_->Free(block, block->size(cache_), block->index(cache_));
     cache_.invalidate(block);
 
     pool = PoolSet::reverse_iterator(pool_.erase(std::next(pool).base()));
 
-    total_free_ -= max_chunk_size_;
+    total_free_ -= block->size(cache_);
     fallback_alloc_count_--;
 
     // If no fall allocation exists, return directly
@@ -322,12 +336,12 @@ void BuddyAllocator::CleanIdleNormalAlloc() {
 
     VLOG(10) << "Return block " << block << " to base allocator.";
 
-    system_allocator_->Free(block, max_chunk_size_, block->index(cache_));
+    system_allocator_->Free(block, block->size(cache_), block->index(cache_));
     cache_.invalidate(block);
 
     pool = PoolSet::reverse_iterator(pool_.erase(std::next(pool).base()));
 
-    total_free_ -= max_chunk_size_;
+    total_free_ -= block->size(cache_);
 
     if (!shall_free_alloc()) return;
   }

paddle/fluid/memory/detail/buddy_allocator.h

@@ -60,7 +60,7 @@ class BuddyAllocator {
   void* SystemAlloc(size_t size);
 
   /*! \brief If existing chunks are not suitable, refill pool */
-  PoolSet::iterator RefillPool();
+  PoolSet::iterator RefillPool(size_t request_bytes);
 
   /**
    *  \brief   Find the suitable chunk from existing pool and split
@@ -89,6 +89,8 @@ class BuddyAllocator {
   size_t min_chunk_size_;  // the minimum size of each chunk
   size_t max_chunk_size_;  // the maximum size of each chunk
 
+  size_t realloc_size_ = 0;  // the size of re-allocated chunk
+
  private:
   /**
    * \brief A list of free allocation

paddle/fluid/memory/detail/buddy_allocator_test.cc

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/memory/detail/buddy_allocator.h"
+
+#include <memory>
+
+#include "gflags/gflags.h"
+#include "gtest/gtest.h"
+#include "paddle/fluid/memory/detail/system_allocator.h"
+#include "paddle/fluid/platform/gpu_info.h"
+
+#ifdef PADDLE_WITH_CUDA
+DECLARE_double(fraction_of_gpu_memory_to_use);
+DECLARE_uint64(initial_gpu_memory_in_mb);
+DECLARE_uint64(reallocate_gpu_memory_in_mb);
+#endif
+
+namespace paddle {
+namespace memory {
+namespace detail {
+
+constexpr static int test_gpu_id = 0;
+
+void TestBuddyAllocator(BuddyAllocator* allocator, size_t size_bytes) {
+  bool freed = false;
+  size_t used_bytes = allocator->Used();
+
+  if (size_bytes > 0) {
+    void* p = allocator->Alloc(size_bytes);
+
+    EXPECT_NE(p, nullptr);
+#ifdef PADDLE_WITH_CUDA
+    if (size_bytes < platform::GpuMaxChunkSize()) {
+#else
+    if (size_bytes < platform::CpuMaxChunkSize()) {
+#endif
+      // Not allocate from SystemAllocator
+      EXPECT_GE(allocator->Used(), used_bytes + size_bytes);
+    } else {
+      // Allocate from SystemAllocator doesn't count in Used()
+      EXPECT_EQ(allocator->Used(), used_bytes);
+    }
+
+    int* intp = static_cast<int*>(p);
+    std::shared_ptr<int> ptr(intp, [&](void* p) {
+      allocator->Free(intp);
+      freed = true;
+    });
+  } else {
+    freed = true;
+  }
+
+  EXPECT_EQ(used_bytes, allocator->Used());
+  EXPECT_TRUE(freed);
+}
+
+#ifdef PADDLE_WITH_CUDA
+TEST(BuddyAllocator, GpuFraction) {
+  FLAGS_fraction_of_gpu_memory_to_use = 0.01;
+
+  BuddyAllocator buddy_allocator(
+      std::unique_ptr<SystemAllocator>(new GPUAllocator(test_gpu_id)),
+      platform::GpuMinChunkSize(), platform::GpuMaxChunkSize());
+
+  TestBuddyAllocator(&buddy_allocator, 10);
+  TestBuddyAllocator(&buddy_allocator, 10 << 10);
+  TestBuddyAllocator(&buddy_allocator, 10 << 20);
+  TestBuddyAllocator(&buddy_allocator, 2 * static_cast<size_t>(1 << 30));
+}
+
+TEST(BuddyAllocator, InitRealloc) {
+  FLAGS_initial_gpu_memory_in_mb = 100;
+  FLAGS_reallocate_gpu_memory_in_mb = 50;
+
+  EXPECT_EQ(platform::GpuMaxChunkSize(), static_cast<size_t>(100 << 20));
+
+  BuddyAllocator buddy_allocator(
+      std::unique_ptr<SystemAllocator>(new GPUAllocator(test_gpu_id)),
+      platform::GpuMinChunkSize(), platform::GpuMaxChunkSize());
+
+  // Less then initial size and reallocate size
+  TestBuddyAllocator(&buddy_allocator, 10 << 20);
+  // Between initial size and reallocate size and not exceed pool
+  TestBuddyAllocator(&buddy_allocator, 80 << 20);
+  // Less then reallocate size and exceed pool
+  TestBuddyAllocator(&buddy_allocator, 40 << 20);
+  // Greater then reallocate size and exceed pool
+  TestBuddyAllocator(&buddy_allocator, 80 << 20);
+  // Greater then initial size and reallocate size
+  TestBuddyAllocator(&buddy_allocator, 2 * static_cast<size_t>(1 << 30));
+}
+
+TEST(BuddyAllocator, ReallocSizeGreaterThanInit) {
+  FLAGS_initial_gpu_memory_in_mb = 5;
+  FLAGS_reallocate_gpu_memory_in_mb = 10;
+
+  EXPECT_EQ(platform::GpuMaxChunkSize(), static_cast<size_t>(10 << 20));
+
+  BuddyAllocator buddy_allocator(
+      std::unique_ptr<SystemAllocator>(new GPUAllocator(test_gpu_id)),
+      platform::GpuMinChunkSize(), platform::GpuMaxChunkSize());
+
+  // Less then initial size and reallocate size
+  TestBuddyAllocator(&buddy_allocator, 1 << 20);
+  // Between initial size and reallocate size and not exceed pool
+  TestBuddyAllocator(&buddy_allocator, 3 << 20);
+  // Less then initial size and exceed pool
+  TestBuddyAllocator(&buddy_allocator, 3 << 20);
+  // Less then reallocate size and not exceed pool (now pool is 15 MB, used 7
+  // MB)
+  TestBuddyAllocator(&buddy_allocator, 7 << 20);
+  // Less then reallocate size and exceed pool
+  TestBuddyAllocator(&buddy_allocator, 8 << 20);
+  // Greater then initial size and reallocate size
+  TestBuddyAllocator(&buddy_allocator, 2 * static_cast<size_t>(1 << 30));
+}
+#endif
+
+}  // namespace detail
+}  // namespace memory
+}  // namespace paddle

paddle/fluid/memory/detail/system_allocator.cc

@@ -32,6 +32,9 @@ limitations under the License. */
 
 DECLARE_bool(use_pinned_memory);
 DECLARE_double(fraction_of_gpu_memory_to_use);
+DECLARE_uint64(initial_gpu_memory_in_mb);
+DECLARE_uint64(reallocate_gpu_memory_in_mb);
+
 namespace paddle {
 namespace memory {
 namespace detail {
@@ -119,11 +122,18 @@ void* GPUAllocator::Alloc(size_t* index, size_t size) {
     gpu_alloc_size_ += size;
     return p;
   } else {
-    LOG(WARNING)
-        << "Cannot malloc " << size / 1024.0 / 1024.0
-        << " MB GPU memory. Please shrink FLAGS_fraction_of_gpu_memory_to_use "
-           "environment variable to a lower value. Current value is "
-        << FLAGS_fraction_of_gpu_memory_to_use;
+    LOG(WARNING) << "Cannot malloc " << size / 1024.0 / 1024.0
+                 << " MB GPU memory. Please shrink "
+                    "FLAGS_fraction_of_gpu_memory_to_use or "
+                    "FLAGS_initial_gpu_memory_in_mb or "
+                    "FLAGS_reallocate_gpu_memory_in_mb"
+                    "environment variable to a lower value. "
+                 << "Current FLAGS_fraction_of_gpu_memory_to_use value is "
+                 << FLAGS_fraction_of_gpu_memory_to_use
+                 << ". Current FLAGS_initial_gpu_memory_in_mb value is "
+                 << FLAGS_initial_gpu_memory_in_mb
+                 << ". Current FLAGS_reallocate_gpu_memory_in_mb value is "
+                 << FLAGS_reallocate_gpu_memory_in_mb;
     return nullptr;
   }
 }

paddle/fluid/operators/fsp_op.cc

+/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/operators/fsp_op.h"
+
+namespace paddle {
+namespace operators {
+
+class FSPOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) of FSPOp should not be null.");
+    PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) of FSPOp should not be null.");
+    PADDLE_ENFORCE(ctx->HasOutput("Out"),
+                   "Output(Out) of FSPOp should not be null.");
+
+    auto x_dims = ctx->GetInputDim("X");
+    auto y_dims = ctx->GetInputDim("Y");
+
+    PADDLE_ENFORCE(
+        x_dims.size() == 4,
+        "The Input(X) must have shape [batch_size, channel, height, width].");
+    PADDLE_ENFORCE(
+        y_dims.size() == 4,
+        "The Input(Y) must have shape [batch_size, channel, height, width].");
+    PADDLE_ENFORCE(
+        (x_dims[2] == y_dims[2]) && (x_dims[3] == y_dims[3]),
+        "The Input(X) and Input(Y) should have the same height and width.");
+
+    ctx->SetOutputDim("Out", {x_dims[0], x_dims[1], y_dims[1]});
+    ctx->ShareLoD("X", "Out");
+  }
+
+ protected:
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext& ctx) const override {
+    framework::LibraryType library_{framework::LibraryType::kPlain};
+    framework::DataLayout layout_ = framework::DataLayout::kAnyLayout;
+    return framework::OpKernelType(ctx.Input<Tensor>("X")->type(),
+                                   ctx.device_context(), layout_, library_);
+  }
+};
+
+class FSPOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("X",
+             "(Tensor) The input of FSP op with shape [batch_size, x_channel, "
+             "height, width]");
+    AddInput("Y",
+             "(Tensor) The input of FSP op with shape"
+             "[batch_size, y_channel, height, width]."
+             "The y_channel can be different with the x_channel of Input(X)"
+             " while the other dimensions must be the same with Input(X)'s.");
+    AddOutput(
+        "Out",
+        "(Tensor) The output of FSP op with shape "
+        "[batch_size, x_channel, y_channel]. The x_channel is the channel "
+        "of Input(X) and the y_channel is the channel of Input(Y).");
+    AddComment(R"DOC(
+    This op is used to calculate the flow of solution procedure (FSP) matrix of two feature maps.
+    Given feature map x with shape [x_channel, h, w] and feature map y with shape
+    [y_channel, h, w], we can get the fsp matrix of x and y in two steps:
+
+        step 1: reshape x into matrix with shape [x_channel, h * w] and reshape and
+                transpose y into matrix with shape [h * w, y_channel]
+        step 2: multiply x and y to get fsp matrix with shape [x_channel, y_channel]
+
+    The output is a batch of fsp matrices.
+    )DOC");
+  }
+};
+
+class FSPOpGrad : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext* ctx) const override {
+    PADDLE_ENFORCE(ctx->HasInput("X"), "Input(X) should not be null");
+    PADDLE_ENFORCE(ctx->HasInput("Y"), "Input(Y) should not be null");
+    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
+                   "Input(Out@GRAD) should not be null");
+    auto x_dims = ctx->GetInputDim("X");
+    auto y_dims = ctx->GetInputDim("Y");
+    auto x_grad_name = framework::GradVarName("X");
+    auto y_grad_name = framework::GradVarName("Y");
+    if (ctx->HasOutput(x_grad_name)) {
+      ctx->SetOutputDim(x_grad_name, x_dims);
+    }
+    if (ctx->HasOutput(y_grad_name)) {
+      ctx->SetOutputDim(y_grad_name, y_dims);
+    }
+  }
+
+  framework::OpKernelType GetExpectedKernelType(
+      const framework::ExecutionContext& ctx) const override {
+    return framework::OpKernelType(
+        ctx.Input<framework::LoDTensor>(framework::GradVarName("Out"))->type(),
+        ctx.device_context());
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OPERATOR(fsp, ops::FSPOp, ops::FSPOpMaker,
+                  paddle::framework::DefaultGradOpDescMaker<true>);
+REGISTER_OPERATOR(fsp_grad, ops::FSPOpGrad);
+REGISTER_OP_CPU_KERNEL(
+    fsp, ops::FSPOpKernel<paddle::platform::CPUDeviceContext, float>,
+    ops::FSPOpKernel<paddle::platform::CPUDeviceContext, double>);
+REGISTER_OP_CPU_KERNEL(
+    fsp_grad, ops::FSPGradOpKernel<paddle::platform::CPUDeviceContext, float>,
+    ops::FSPGradOpKernel<paddle::platform::CPUDeviceContext, double>);

paddle/fluid/operators/fsp_op.cu

+/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/fsp_op.h"
+
+namespace ops = paddle::operators;
+namespace plat = paddle::platform;
+REGISTER_OP_CUDA_KERNEL(fsp, ops::FSPOpKernel<plat::CUDADeviceContext, float>,
+                        ops::FSPOpKernel<plat::CUDADeviceContext, double>);
+REGISTER_OP_CUDA_KERNEL(fsp_grad,
+                        ops::FSPGradOpKernel<plat::CUDADeviceContext, float>,
+                        ops::FSPGradOpKernel<plat::CUDADeviceContext, double>);

paddle/fluid/operators/fsp_op.h

+/* Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/math/blas.h"
+#include "paddle/fluid/operators/math/math_function.h"
+
+namespace paddle {
+namespace operators {
+
+using Tensor = framework::Tensor;
+
+template <typename DeviceContext, typename T>
+class FSPOpKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    auto* x = context.Input<Tensor>("X");
+    auto* y = context.Input<Tensor>("Y");
+    auto* output = context.Output<Tensor>("Out");
+    output->mutable_data<T>(context.GetPlace());
+    auto x_dims = x->dims();
+    auto y_dims = y->dims();
+
+    auto batch_size = x_dims[0];
+    auto x_channel = x_dims[1];
+    auto y_channel = y_dims[1];
+    auto height = x_dims[2];
+    auto width = x_dims[3];
+
+    auto blas = math::GetBlas<DeviceContext, T>(context);
+
+    math::MatDescriptor x_mat_desc;
+    x_mat_desc.height_ = x_channel;
+    x_mat_desc.width_ = height * width;
+    x_mat_desc.batch_size_ = batch_size;
+    x_mat_desc.stride_ = x_channel * height * width;
+
+    math::MatDescriptor y_mat_desc;
+    y_mat_desc.height_ = height * width;
+    y_mat_desc.width_ = y_channel;
+    y_mat_desc.batch_size_ = batch_size;
+    y_mat_desc.stride_ = y_channel * height * width;
+    y_mat_desc.trans_ = true;
+
+    blas.MatMul(*x, x_mat_desc, *y, y_mat_desc,
+                static_cast<T>(1.0 / (height * width)), output,
+                static_cast<T>(0.0));
+  }
+};
+
+template <typename DeviceContext, typename T>
+class FSPGradOpKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    auto* d_x = context.Output<Tensor>(framework::GradVarName("X"));
+    auto* d_y = context.Output<Tensor>(framework::GradVarName("Y"));
+    if (d_x == nullptr && d_y == nullptr) {
+      return;
+    }
+    auto* d_out = context.Input<Tensor>(framework::GradVarName("Out"));
+    auto d_out_dims = d_out->dims();
+    auto batch_size = d_out_dims[0];
+    auto x_channel = d_out_dims[1];
+    auto y_channel = d_out_dims[2];
+    int64_t h = 0;
+    int64_t w = 0;
+
+    auto blas = math::GetBlas<DeviceContext, T>(context);
+    math::SetConstant<DeviceContext, T> set_zero;
+    if (d_x != nullptr) {
+      d_x->mutable_data<T>(context.GetPlace());
+      set_zero(context.template device_context<DeviceContext>(), d_x,
+               static_cast<T>(0));
+      auto* y = context.Input<Tensor>("Y");
+      auto y_dims = y->dims();
+      h = y_dims[2];
+      w = y_dims[3];
+
+      math::MatDescriptor d_out_mat_desc;
+      d_out_mat_desc.height_ = x_channel;
+      d_out_mat_desc.width_ = y_channel;
+      d_out_mat_desc.batch_size_ = batch_size;
+      d_out_mat_desc.stride_ = x_channel * y_channel;
+
+      math::MatDescriptor y_mat_desc;
+      y_mat_desc.height_ = y_channel;
+      y_mat_desc.width_ = h * w;
+      y_mat_desc.batch_size_ = batch_size;
+      y_mat_desc.stride_ = y_channel * h * w;
+
+      blas.MatMul(*d_out, d_out_mat_desc, *y, y_mat_desc,
+                  static_cast<T>(1.0 / (h * w)), d_x, static_cast<T>(0.0));
+    }
+
+    if (d_y != nullptr) {
+      d_y->mutable_data<T>(context.GetPlace());
+      set_zero(context.template device_context<DeviceContext>(), d_y,
+               static_cast<T>(0));
+      auto* x = context.Input<Tensor>("X");
+      auto x_dims = x->dims();
+      h = x_dims[2];
+      w = x_dims[3];
+
+      math::MatDescriptor d_out_mat_desc;
+      d_out_mat_desc.height_ = y_channel;
+      d_out_mat_desc.width_ = x_channel;
+      d_out_mat_desc.batch_size_ = batch_size;
+      d_out_mat_desc.stride_ = x_channel * y_channel;
+      d_out_mat_desc.trans_ = true;
+
+      math::MatDescriptor x_mat_desc;
+      x_mat_desc.height_ = x_channel;
+      x_mat_desc.width_ = h * w;
+      x_mat_desc.batch_size_ = batch_size;
+      x_mat_desc.stride_ = x_channel * h * w;
+
+      blas.MatMul(*d_out, d_out_mat_desc, *x, x_mat_desc,
+                  static_cast<T>(1.0 / (h * w)), d_y, static_cast<T>(0.0));
+    }
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/operators/range_op.cc

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/operators/range_op.h"
+
+namespace paddle {
+namespace operators {
+
+class RangeOp : public framework::OperatorWithKernel {
+ public:
+  using framework::OperatorWithKernel::OperatorWithKernel;
+
+  void InferShape(framework::InferShapeContext *ctx) const override {
+    if (ctx->HasInput("Start")) {
+      auto s_dims = ctx->GetInputDim("Start");
+      PADDLE_ENFORCE((s_dims.size() == 1) && (s_dims[0] == 1),
+                     "The shape of Input(Start) should be [1].");
+    }
+    if (ctx->HasInput("End")) {
+      auto e_dims = ctx->GetInputDim("End");
+      PADDLE_ENFORCE((e_dims.size() == 1) && (e_dims[0] == 1),
+                     "The shape of Input(End) should be [1].");
+    }
+    if (ctx->HasInput("Step")) {
+      auto step_dims = ctx->GetInputDim("Step");
+      PADDLE_ENFORCE((step_dims.size() == 1) && (step_dims[0] == 1),
+                     "The shape of Input(Step) should be [1].");
+    }
+    ctx->SetOutputDim("Out", {-1});
+  }
+};
+
+class RangeOpMaker : public framework::OpProtoAndCheckerMaker {
+ public:
+  void Make() override {
+    AddInput("Start",
+             "Start of interval. The interval includes this value. It is a "
+             "tensor with shape=[1].");
+    AddInput("End",
+             "End of interval. The interval does not include this value, "
+             "except in some cases where step is not an integer and floating "
+             "point round-off affects the length of out. It is a tensor with "
+             "shape=[1].");
+    AddInput("Step", "Spacing between values. It is a tensor with shape=[1].");
+    AddOutput("Out", "A sequence of numbers.");
+    AddComment(R"DOC(
+    Return evenly spaced values within a given interval. Values are generated within the half-open interval [start, stop) (in other words, the interval including start but excluding stop). Like arange function of numpy.
+)DOC");
+  }
+};
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OP_WITHOUT_GRADIENT(range, ops::RangeOp, ops::RangeOpMaker);
+REGISTER_OP_CPU_KERNEL(range, ops::CPURangeKernel<int>,
+                       ops::CPURangeKernel<float>, ops::CPURangeKernel<double>,
+                       ops::CPURangeKernel<int64_t>);

paddle/fluid/operators/range_op.cu

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/range_op.h"
+#include "paddle/fluid/platform/cuda_primitives.h"
+
+namespace paddle {
+namespace operators {
+
+#define CUDA_1D_KERNEL_LOOP(i, n)                              \
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < (n); \
+       i += blockDim.x * gridDim.x)
+
+template <typename T>
+__global__ void RangeKernel(T start, T step, int64_t size, T* out) {
+  CUDA_1D_KERNEL_LOOP(index, size) { out[index] = start + step * index; }
+}
+
+template <typename T>
+class CUDARangeKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    auto* start_t = context.Input<framework::Tensor>("Start");
+    auto* end_t = context.Input<framework::Tensor>("End");
+    auto* step_t = context.Input<framework::Tensor>("Step");
+    auto* out = context.Output<framework::Tensor>("Out");
+
+    framework::Tensor n;
+    framework::TensorCopy(*start_t, platform::CPUPlace(), &n);
+    T start = n.data<T>()[0];
+    framework::TensorCopy(*end_t, platform::CPUPlace(), &n);
+    T end = n.data<T>()[0];
+    framework::TensorCopy(*step_t, platform::CPUPlace(), &n);
+    T step = n.data<T>()[0];
+
+    int64_t size = 0;
+    GetSize(start, end, step, &size);
+    out->Resize(framework::make_ddim({size}));
+    T* out_data = out->mutable_data<T>(context.GetPlace());
+
+    auto stream = context.cuda_device_context().stream();
+    int block = 512;
+    int grid = (size + block - 1) / block;
+    RangeKernel<T><<<grid, block, 0, stream>>>(start, step, size, out_data);
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle
+
+namespace ops = paddle::operators;
+REGISTER_OP_CUDA_KERNEL(range, ops::CUDARangeKernel<int>,
+                        ops::CUDARangeKernel<int64_t>,
+                        ops::CUDARangeKernel<float>,
+                        ops::CUDARangeKernel<double>);

paddle/fluid/operators/range_op.h

+/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#pragma once
+#include <functional>
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/operators/math/math_function.h"
+
+namespace paddle {
+namespace operators {
+
+template <typename T>
+void GetSize(T start, T end, T step, int64_t* size) {
+  PADDLE_ENFORCE(!std::equal_to<T>()(step, 0),
+                 "The step of range op should not be 0.");
+  PADDLE_ENFORCE(((start < end) && (step > 0)) || ((start > end) && (step < 0)),
+                 "The step should be greater than 0 while start < end. And the "
+                 "step should be less than 0 while start > end.");
+  *size = std::is_integral<T>::value
+              ? ((std::abs(end - start) + std::abs(step) - 1) / std::abs(step))
+              : std::ceil(std::abs((end - start) / step));
+}
+
+template <typename T>
+class CPURangeKernel : public framework::OpKernel<T> {
+ public:
+  void Compute(const framework::ExecutionContext& context) const override {
+    T start = context.Input<framework::Tensor>("Start")->data<T>()[0];
+    T end = context.Input<framework::Tensor>("End")->data<T>()[0];
+    T step = context.Input<framework::Tensor>("Step")->data<T>()[0];
+    auto* out = context.Output<framework::Tensor>("Out");
+    int64_t size = 0;
+    GetSize(start, end, step, &size);
+    out->Resize(framework::make_ddim({size}));
+    T* out_data = out->mutable_data<T>(context.GetPlace());
+    T value = start;
+    for (int64_t i = 0; i < size; ++i) {
+      out_data[i] = value;
+      value += step;
+    }
+  }
+};
+
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/operators/reader/CMakeLists.txt

@@ -17,7 +17,9 @@ function(reader_library TARGET_NAME)
         PARENT_SCOPE)
 endfunction()
 
+cc_library(py_reader SRCS py_reader.cc DEPS reader)
 cc_library(buffered_reader SRCS buffered_reader.cc DEPS reader simple_threadpool)
+
 reader_library(open_files_op SRCS open_files_op.cc DEPS buffered_reader)
 reader_library(create_random_data_generator_op SRCS create_random_data_generator_op.cc)
 reader_library(create_shuffle_reader_op SRCS create_shuffle_reader_op.cc)
@@ -26,7 +28,7 @@ reader_library(create_recordio_file_reader_op SRCS create_recordio_file_reader_o
 reader_library(create_double_buffer_reader_op SRCS create_double_buffer_reader_op.cc DEPS buffered_reader)
 reader_library(create_multi_pass_reader_op SRCS create_multi_pass_reader_op.cc)
 reader_library(create_custom_reader_op SRCS create_custom_reader_op.cc)
-reader_library(create_py_reader_op SRCS create_py_reader_op.cc)
+reader_library(create_py_reader_op SRCS create_py_reader_op.cc DEPS py_reader)
 
 if (NOT WIN32 AND NOT ON_INFER)
     cc_library(ctr_reader SRCS ctr_reader.cc DEPS gzstream reader zlib)
@@ -38,7 +40,7 @@ cc_test(reader_blocking_queue_test SRCS reader_blocking_queue_test.cc)
 # Export local libraries to parent
 # set(READER_LIBRARY ${LOCAL_READER_LIBS} PARENT_SCOPE)
 
-op_library(read_op)
+op_library(read_op DEPS py_reader buffered_reader)
 
 foreach(src ${LOCAL_READER_LIBS})
     set(OP_LIBRARY ${src} ${OP_LIBRARY} CACHE INTERNAL "op libs")

paddle/fluid/operators/reader/blocking_queue.h

@@ -16,6 +16,7 @@
 
 #include <condition_variable>  // NOLINT
 #include <deque>
+#include <utility>
 
 #include "paddle/fluid/platform/enforce.h"
 
@@ -34,7 +35,7 @@ class BlockingQueue {
   explicit BlockingQueue(size_t capacity, bool speed_test_mode = false)
       : capacity_(capacity), speed_test_mode_(speed_test_mode), closed_(false) {
     PADDLE_ENFORCE_GT(
-        capacity_, 0,
+        capacity_, static_cast<size_t>(0),
         "The capacity of a reader::BlockingQueue must be greater than 0.");
   }
 

paddle/fluid/operators/reader/buffered_reader.cc

@@ -30,8 +30,10 @@ BufferedReader::~BufferedReader() {
 #ifdef PADDLE_WITH_CUDA
   if (platform::is_gpu_place(place_)) {
     platform::SetDeviceId(boost::get<platform::CUDAPlace>(place_).device);
-    PADDLE_ENFORCE(cudaStreamDestroy(stream));
-    for (auto &event : events) PADDLE_ENFORCE(cudaEventDestroy(event));
+    PADDLE_ENFORCE(cudaStreamDestroy(stream_));
+    for (auto &event : events_) {
+      PADDLE_ENFORCE(cudaEventDestroy(event));
+    }
   }
 #endif
 }
@@ -46,15 +48,15 @@ BufferedReader::BufferedReader(
 #ifdef PADDLE_WITH_CUDA
   if (platform::is_gpu_place(place_)) {
     platform::SetDeviceId(boost::get<platform::CUDAPlace>(place_).device);
-    compute_stream =
+    compute_stream_ =
         ((platform::CUDADeviceContext *)(platform::DeviceContextPool::Instance()
                                              .Get(place_)))
             ->stream();
-    events.resize(buffer_size);
-    for (auto &event : events) {
+    events_.resize(buffer_size);
+    for (auto &event : events_) {
       PADDLE_ENFORCE(cudaEventCreateWithFlags(&event, cudaEventDisableTiming));
     }
-    PADDLE_ENFORCE(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking));
+    PADDLE_ENFORCE(cudaStreamCreateWithFlags(&stream_, cudaStreamNonBlocking));
   }
 #endif
   cpu_buffer_.resize(buffer_size);
@@ -73,7 +75,7 @@ void BufferedReader::ReadAsync(size_t i) {
 #ifdef PADDLE_WITH_CUDA
   if (platform::is_gpu_place(place_)) {
     platform::SetDeviceId(boost::get<platform::CUDAPlace>(place_).device);
-    PADDLE_ENFORCE(cudaEventRecord(events[i], compute_stream));
+    PADDLE_ENFORCE(cudaEventRecord(events_[i], compute_stream_));
   }
 #endif
   position_.emplace(thread_pool_.enqueue([this, i]() -> size_t {
@@ -91,7 +93,7 @@ void BufferedReader::ReadAsync(size_t i) {
     // commands from different streams cannot run concurrently.
     if (platform::is_gpu_place(place_)) {
       platform::SetDeviceId(boost::get<platform::CUDAPlace>(place_).device);
-      PADDLE_ENFORCE(cudaStreamWaitEvent(stream, events[i], 0));
+      PADDLE_ENFORCE(cudaStreamWaitEvent(stream_, events_[i], 0));
       TensorVec &gpu = gpu_buffer_[i];
       gpu.resize(cpu.size());
       platform::RecordEvent record_event("BufferedReader:MemoryCopy");
@@ -106,12 +108,14 @@ void BufferedReader::ReadAsync(size_t i) {
         if (platform::is_cuda_pinned_place(cpu_place)) {
           memory::Copy(boost::get<platform::CUDAPlace>(place_), gpu_ptr,
                        boost::get<platform::CUDAPinnedPlace>(cpu_place),
-                       cpu_ptr, size, stream);
+                       cpu_ptr, size, stream_);
         } else if ((platform::is_gpu_place(cpu_place))) {
           memory::Copy(boost::get<platform::CUDAPlace>(place_), gpu_ptr,
                        boost::get<platform::CUDAPlace>(cpu_place), cpu_ptr,
-                       size, stream);
+                       size, stream_);
         } else {
+          // if cpu place is not pinned, async copy is slower than sync copy,
+          // so we use sync copy instead.
           // TODO(zcd): The default stream should not be used here.
           memory::Copy(boost::get<platform::CUDAPlace>(place_), gpu_ptr,
                        boost::get<platform::CPUPlace>(cpu_place), cpu_ptr, size,
@@ -119,7 +123,7 @@ void BufferedReader::ReadAsync(size_t i) {
         }
         gpu[i].set_lod(cpu[i].lod());
       }
-      PADDLE_ENFORCE(cudaStreamSynchronize(stream));
+      PADDLE_ENFORCE(cudaStreamSynchronize(stream_));
     }
 #endif
     return i;

paddle/fluid/operators/reader/buffered_reader.h

@@ -15,6 +15,7 @@
 #pragma once
 
 #include <list>
+#include <memory>
 #include <queue>
 #include <vector>
 #include "ThreadPool.h"
@@ -63,9 +64,9 @@ class BufferedReader : public framework::DecoratedReader {
   std::vector<TensorVec> gpu_buffer_;
   size_t prev_pos_{-1UL};
 #ifdef PADDLE_WITH_CUDA
-  cudaStream_t stream;
-  cudaStream_t compute_stream;
-  std::vector<cudaEvent_t> events;
+  cudaStream_t stream_;
+  cudaStream_t compute_stream_;
+  std::vector<cudaEvent_t> events_;
 #endif
 };
 

paddle/fluid/operators/reader/create_py_reader_op.cc

@@ -12,37 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#include "paddle/fluid/operators/reader/lod_tensor_blocking_queue.h"
+#include "paddle/fluid/operators/reader/py_reader.h"
 #include "paddle/fluid/operators/reader/reader_op_registry.h"
 
 namespace paddle {
 namespace operators {
 namespace reader {
 
-class PyReader : public framework::FileReader {
- public:
-  explicit PyReader(const std::shared_ptr<LoDTensorBlockingQueue>& queue)
-      : framework::FileReader() {
-    PADDLE_ENFORCE(queue != nullptr, "LoDTensorBlockingQueue must not be null");
-    queue_ = queue;
-  }
-
-  void ReadNext(std::vector<framework::LoDTensor>* out) override {
-    bool success;
-    *out = queue_->Pop(&success);
-    if (!success) out->clear();
-  }
-
-  ~PyReader() { queue_->Close(); }
-
-  void Shutdown() override { queue_->Close(); }
-
-  void Start() override { queue_->ReOpen(); }
-
- private:
-  std::shared_ptr<LoDTensorBlockingQueue> queue_;
-};
-
 class CreatePyReaderOp : public framework::OperatorBase {
  public:
   using framework::OperatorBase::OperatorBase;

paddle/fluid/operators/reader/py_reader.cc

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/operators/reader/py_reader.h"
+#include <memory>
+
+namespace paddle {
+namespace operators {
+namespace reader {
+
+PyReader::PyReader(const std::shared_ptr<LoDTensorBlockingQueue>& queue)
+    : framework::FileReader() {
+  PADDLE_ENFORCE(queue != nullptr, "LoDTensorBlockingQueue must not be null");
+  queue_ = queue;
+}
+
+void PyReader::ReadNext(std::vector<framework::LoDTensor>* out) {
+  bool success;
+  *out = queue_->Pop(&success);
+  if (!success) out->clear();
+}
+
+PyReader::~PyReader() { queue_->Close(); }
+
+void PyReader::Shutdown() { queue_->Close(); }
+
+void PyReader::Start() { queue_->ReOpen(); }
+
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/operators/reader/py_reader.h

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <atomic>
+#include <memory>
+#include <vector>
+#include "paddle/fluid/framework/reader.h"
+#include "paddle/fluid/operators/reader/lod_tensor_blocking_queue.h"
+
+namespace paddle {
+namespace operators {
+namespace reader {
+
+class PyReader : public framework::FileReader {
+ public:
+  explicit PyReader(const std::shared_ptr<LoDTensorBlockingQueue>& queue);
+
+  void ReadNext(std::vector<framework::LoDTensor>* out) override;
+
+  ~PyReader();
+
+  void Shutdown() override;
+
+  void Start() override;
+
+ private:
+  std::shared_ptr<LoDTensorBlockingQueue> queue_;
+};
+
+}  // namespace reader
+}  // namespace operators
+}  // namespace paddle

paddle/fluid/platform/gpu_info.cc

@@ -13,7 +13,6 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #include "paddle/fluid/platform/gpu_info.h"
-
 #include <algorithm>
 #include <cstdlib>
 #include <string>
@@ -31,6 +30,8 @@ constexpr static float fraction_of_gpu_memory_to_use = 0.92f;
 constexpr static float fraction_of_gpu_memory_to_use = 0.5f;
 #endif
 
+constexpr static float fraction_reserve_gpu_memory = 0.05f;
+
 DEFINE_double(fraction_of_gpu_memory_to_use, fraction_of_gpu_memory_to_use,
               "Allocate a trunk of gpu memory that is this fraction of the "
               "total gpu memory size. Future memory usage will be allocated "
@@ -38,6 +39,24 @@ DEFINE_double(fraction_of_gpu_memory_to_use, fraction_of_gpu_memory_to_use,
               "additional trunks of the same size will be requested from gpu "
               "until the gpu has no memory left for another trunk.");
 
+DEFINE_uint64(
+    initial_gpu_memory_in_mb, 0ul,
+    "Allocate a trunk of gpu memory whose byte size is specified by "
+    "the flag. Future memory usage will be allocated from the "
+    "truck. If the trunk doesn't have enough gpu memory, additional "
+    "trunks of the gpu memory will be requested from gpu with size "
+    "specified by FLAGS_reallocate_gpu_memory_in_mb until the gpu has "
+    "no memory left for the additional trunk. Note: if you set this "
+    "flag, the memory size set by "
+    "FLAGS_fraction_of_gpu_memory_to_use will be overrided by this "
+    "flag. If you don't set this flag, PaddlePaddle will use "
+    "FLAGS_fraction_of_gpu_memory_to_use to allocate gpu memory");
+
+DEFINE_uint64(reallocate_gpu_memory_in_mb, 0ul,
+              "If this flag is set, Paddle will reallocate the gpu memory with "
+              "size specified by this flag. Else Paddle will reallocate by "
+              "FLAGS_fraction_of_gpu_memory_to_use");
+
 DEFINE_bool(
     enable_cublas_tensor_op_math, false,
     "The enable_cublas_tensor_op_math indicate whether to use Tensor Core, "
@@ -180,13 +199,43 @@ void GpuMemoryUsage(size_t *available, size_t *total) {
 }
 
 size_t GpuMaxAllocSize() {
+  return std::max(GpuInitAllocSize(), GpuReallocSize());
+}
+
+size_t GpuInitAllocSize() {
+  if (FLAGS_initial_gpu_memory_in_mb > 0ul) {
+    // Initial memory will be allocated by FLAGS_initial_gpu_memory_in_mb
+    return static_cast<size_t>(FLAGS_initial_gpu_memory_in_mb << 20);
+  }
+
+  // FLAGS_initial_gpu_memory_in_mb is 0, initial memory will be allocated by
+  // fraction
   size_t total = 0;
   size_t available = 0;
 
   GpuMemoryUsage(&available, &total);
+  size_t reserving = static_cast<size_t>(fraction_reserve_gpu_memory * total);
 
-  // Reserve the rest for page tables, etc.
-  return static_cast<size_t>(total * FLAGS_fraction_of_gpu_memory_to_use);
+  return static_cast<size_t>((total - reserving) *
+                             FLAGS_fraction_of_gpu_memory_to_use);
+}
+
+size_t GpuReallocSize() {
+  if (FLAGS_reallocate_gpu_memory_in_mb > 0ul) {
+    // Additional memory will be allocated by FLAGS_reallocate_gpu_memory_in_mb
+    return static_cast<size_t>(FLAGS_reallocate_gpu_memory_in_mb << 20);
+  }
+
+  // FLAGS_reallocate_gpu_memory_in_mb is 0, additional memory will be allocated
+  // by fraction
+  size_t total = 0;
+  size_t available = 0;
+
+  GpuMemoryUsage(&available, &total);
+  size_t reserving = static_cast<size_t>(fraction_reserve_gpu_memory * total);
+
+  return static_cast<size_t>((total - reserving) *
+                             FLAGS_fraction_of_gpu_memory_to_use);
 }
 
 size_t GpuMinChunkSize() {
@@ -201,16 +250,13 @@ size_t GpuMaxChunkSize() {
   GpuMemoryUsage(&available, &total);
   VLOG(10) << "GPU Usage " << available / 1024 / 1024 << "M/"
            << total / 1024 / 1024 << "M";
-  size_t reserving = static_cast<size_t>(0.05 * total);
+  size_t reserving = static_cast<size_t>(fraction_reserve_gpu_memory * total);
   // If available less than minimum chunk size, no usable memory exists.
   available =
       std::min(std::max(available, GpuMinChunkSize()) - GpuMinChunkSize(),
                total - reserving);
 
-  // Reserving the rest memory for page tables, etc.
-
-  size_t allocating = static_cast<size_t>(FLAGS_fraction_of_gpu_memory_to_use *
-                                          (total - reserving));
+  size_t allocating = GpuMaxAllocSize();
 
   PADDLE_ENFORCE_LE(allocating, available,
                     "Insufficient GPU memory to allocation.");

paddle/fluid/platform/gpu_info.h

@@ -60,6 +60,12 @@ void GpuMemoryUsage(size_t *available, size_t *total);
 //! Get the maximum allocation size of current GPU device.
 size_t GpuMaxAllocSize();
 
+//! Get the initial allocation size of current GPU device.
+size_t GpuInitAllocSize();
+
+//! Get the re-allocation size of current GPU device.
+size_t GpuReallocSize();
+
 //! Get the minimum chunk size for GPU buddy allocator.
 size_t GpuMinChunkSize();
 

paddle/fluid/pybind/CMakeLists.txt

@@ -5,7 +5,7 @@ set(PYBIND_DEPS pybind python proto_desc memory executor async_executor prune
 if(WITH_PYTHON)
   list(APPEND PYBIND_DEPS py_func_op)
 endif()
-set(PYBIND_SRCS pybind.cc exception.cc protobuf.cc const_value.cc recordio.cc async_executor_py.cc imperative.cc ir.cc inference_api.cc)
+set(PYBIND_SRCS pybind.cc exception.cc protobuf.cc const_value.cc recordio.cc reader_py.cc async_executor_py.cc imperative.cc ir.cc inference_api.cc)
 
 if(WITH_PYTHON)
   if(WITH_AMD_GPU)

paddle/fluid/pybind/pybind.cc

@@ -55,6 +55,7 @@ limitations under the License. */
 #include "paddle/fluid/pybind/ir.h"
 #include "paddle/fluid/pybind/protobuf.h"
 #include "paddle/fluid/pybind/pybind.h"  // NOLINT
+#include "paddle/fluid/pybind/reader_py.h"
 #include "paddle/fluid/pybind/recordio.h"
 #include "paddle/fluid/pybind/tensor_py.h"
 
@@ -128,6 +129,11 @@ static inline bool IsSamePlace(const PlaceType1 &p1, const PlaceType2 &p2) {
   return paddle::platform::Place(p1) == paddle::platform::Place(p2);
 }
 
+template <typename PlaceType>
+static inline int PlaceIndex(const PlaceType &p) {
+  return static_cast<int>(paddle::platform::Place(p).which());
+}
+
 PYBIND11_MODULE(core, m) {
   // Not used, just make sure cpu_info.cc is linked.
   paddle::platform::CpuTotalPhysicalMemory();
@@ -531,6 +537,7 @@ PYBIND11_MODULE(core, m) {
 
 All parameter, weight, gradient are variables in Paddle.
 )DOC")
+      .def(py::init<>())
       .def("is_int", [](const Variable &var) { return var.IsType<int>(); })
       .def("set_int",
            [](Variable &var, int val) -> void { *var.GetMutable<int>() = val; })
@@ -572,14 +579,13 @@ All parameter, weight, gradient are variables in Paddle.
            },
            py::return_value_policy::reference);
 
-  py::class_<framework::ReaderHolder>(m, "Reader", "")
-      .def("start", &framework::ReaderHolder::Start)
-      .def("reset", &framework::ReaderHolder::ResetAll);
+  BindReader(&m);
 
   using LoDTensorBlockingQueue =
       ::paddle::operators::reader::LoDTensorBlockingQueue;
   using LoDTensorBlockingQueueHolder =
       ::paddle::operators::reader::LoDTensorBlockingQueueHolder;
+
   py::class_<LoDTensorBlockingQueue, std::shared_ptr<LoDTensorBlockingQueue>>(
       m, "LoDTensorBlockingQueue", "")
       .def("push",
@@ -776,6 +782,7 @@ All parameter, weight, gradient are variables in Paddle.
              PADDLE_THROW("Cannot use CUDAPlace in CPU only version");
 #endif
            })
+      .def("_type", &PlaceIndex<platform::CUDAPlace>)
       .def("_equals", &IsSamePlace<platform::CUDAPlace, platform::Place>)
       .def("_equals", &IsSamePlace<platform::CUDAPlace, platform::CUDAPlace>)
       .def("_equals", &IsSamePlace<platform::CUDAPlace, platform::CPUPlace>)
@@ -785,6 +792,7 @@ All parameter, weight, gradient are variables in Paddle.
 
   py::class_<paddle::platform::CPUPlace>(m, "CPUPlace")
       .def(py::init<>())
+      .def("_type", &PlaceIndex<platform::CPUPlace>)
       .def("_equals", &IsSamePlace<platform::CPUPlace, platform::Place>)
       .def("_equals", &IsSamePlace<platform::CPUPlace, platform::CUDAPlace>)
       .def("_equals", &IsSamePlace<platform::CPUPlace, platform::CPUPlace>)
@@ -800,6 +808,7 @@ All parameter, weight, gradient are variables in Paddle.
 #endif
              new (&self) platform::CUDAPinnedPlace();
            })
+      .def("_type", &PlaceIndex<platform::CUDAPinnedPlace>)
       .def("_equals", &IsSamePlace<platform::CUDAPinnedPlace, platform::Place>)
       .def("_equals",
            &IsSamePlace<platform::CUDAPinnedPlace, platform::CUDAPlace>)
@@ -811,16 +820,25 @@ All parameter, weight, gradient are variables in Paddle.
 
   py::class_<platform::Place>(m, "Place")
       .def(py::init<>())
+      .def("_type", &PlaceIndex<platform::Place>)
       .def("_equals", &IsSamePlace<platform::Place, platform::Place>)
       .def("_equals", &IsSamePlace<platform::Place, platform::CUDAPlace>)
       .def("_equals", &IsSamePlace<platform::Place, platform::CPUPlace>)
       .def("_equals", &IsSamePlace<platform::Place, platform::CUDAPinnedPlace>)
       .def("is_gpu_place",
            [](platform::Place &self) { return platform::is_gpu_place(self); })
+      .def("is_cpu_place",
+           [](platform::Place &self) { return platform::is_cpu_place(self); })
+      .def("is_cuda_pinned_place",
+           [](platform::Place &self) {
+             return platform::is_cuda_pinned_place(self);
+           })
       .def("gpu_device_id",
            [](platform::Place &self) {
              return boost::get<platform::CUDAPlace>(self).device;
            })
+      .def("set_place", [](platform::Place &self,
+                           const platform::Place &other) { self = other; })
       .def("set_place",
            [](platform::Place &self, const platform::CPUPlace &cpu_place) {
              self = cpu_place;

paddle/fluid/pybind/reader_py.cc

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/pybind/reader_py.h"
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+#include "paddle/fluid/framework/reader.h"
+#include "paddle/fluid/operators/reader/buffered_reader.h"
+#include "paddle/fluid/operators/reader/py_reader.h"
+#include "paddle/fluid/platform/place.h"
+#include "pybind11/stl.h"
+
+namespace paddle {
+namespace pybind {
+
+class MultiDeviceFeedReader {
+ public:
+  using ResultDictList =
+      std::vector<std::unordered_map<std::string, framework::LoDTensor>>;
+
+  MultiDeviceFeedReader(
+      const std::shared_ptr<operators::reader::LoDTensorBlockingQueue> &queue,
+      const std::vector<std::string> &names,
+      const std::vector<platform::Place> &dst_places, bool use_double_buffer)
+      : queue_(queue),
+        names_(names),
+        pool_(new ::ThreadPool(dst_places.size())) {
+    std::shared_ptr<framework::ReaderBase> reader(
+        new operators::reader::PyReader(queue));
+
+    readers_.reserve(dst_places.size());
+    for (auto &p : dst_places) {
+      auto *holder = new framework::ReaderHolder();
+      if (use_double_buffer) {
+        holder->Reset(
+            framework::MakeDecoratedReader<operators::reader::BufferedReader>(
+                reader, p, 2));
+      } else {
+        if (platform::is_gpu_place(p)) {
+          PADDLE_THROW(
+              "Place cannot be CUDAPlace when use_double_buffer is False");
+        }
+        holder->Reset(reader);
+      }
+      readers_.emplace_back(holder);
+    }
+
+    futures_.resize(dst_places.size());
+    ret_.resize(dst_places.size());
+    ReadAsync();
+  }
+
+  ResultDictList ReadNext() {
+    bool success = WaitFutures();
+
+    if (!success) {
+      return {};
+    }
+
+    ResultDictList result(ret_.size());
+    for (size_t i = 0; i < ret_.size(); ++i) {
+      for (size_t j = 0; j < names_.size(); ++j) {
+        result[i].emplace(names_[j], std::move(ret_[i][j]));
+      }
+    }
+    ReadAsync();
+    return result;
+  }
+
+  void Reset() {
+    Shutdown();
+    Start();
+    ReadAsync();
+  }
+
+  ~MultiDeviceFeedReader() {
+    queue_->Close();
+    pool_.reset();
+  }
+
+ private:
+  bool WaitFutures() {
+    bool success = true;
+    for (auto &f : futures_) {
+      success &= f.get();
+    }
+    return success;
+  }
+
+  void Shutdown() {
+    for (auto &r : readers_) r->Shutdown();
+  }
+
+  void Start() {
+    for (auto &r : readers_) r->Start();
+  }
+
+  void ReadAsync() {
+    for (size_t i = 0; i < readers_.size(); ++i) {
+      futures_[i] = pool_->enqueue([this, i] {
+        readers_[i]->ReadNext(&ret_[i]);
+        return !ret_[i].empty();
+      });
+    }
+  }
+
+  std::shared_ptr<operators::reader::LoDTensorBlockingQueue> queue_;
+  std::vector<std::string> names_;
+  std::unique_ptr<::ThreadPool> pool_;
+
+  std::vector<std::unique_ptr<framework::ReaderHolder>> readers_;
+
+  std::vector<std::future<bool>> futures_;
+  std::vector<std::vector<framework::LoDTensor>> ret_;
+};
+
+namespace py = pybind11;
+
+void BindReader(py::module *module) {
+  auto &m = *module;
+
+  namespace reader = ::paddle::operators::reader;
+
+  py::class_<framework::ReaderHolder>(m, "Reader", "")
+      .def("start", &framework::ReaderHolder::Start)
+      .def("reset", &framework::ReaderHolder::ResetAll);
+
+  py::class_<MultiDeviceFeedReader>(m, "MultiDeviceFeedReader", "")
+      .def("read_next", &MultiDeviceFeedReader::ReadNext,
+           py::call_guard<py::gil_scoped_release>())
+      .def("reset", &MultiDeviceFeedReader::Reset,
+           py::call_guard<py::gil_scoped_release>());
+
+  m.def("create_py_reader",
+        [](const std::shared_ptr<operators::reader::LoDTensorBlockingQueue>
+               &queue,
+           const std::vector<std::string> &names,
+           const std::vector<platform::Place> &dst_places,
+           bool use_double_buffer) {
+          return new MultiDeviceFeedReader(queue, names, dst_places,
+                                           use_double_buffer);
+        },
+        py::return_value_policy::take_ownership);
+}
+
+}  // namespace pybind
+}  // namespace paddle

paddle/fluid/pybind/reader_py.h

+// Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "pybind11/pybind11.h"
+
+namespace paddle {
+namespace pybind {
+
+void BindReader(pybind11::module *module);
+
+}  // namespace pybind
+}  // namespace paddle

paddle/testing/paddle_gtest_main.cc

@@ -41,6 +41,8 @@ int main(int argc, char** argv) {
 
 #if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
   envs.push_back("fraction_of_gpu_memory_to_use");
+  envs.push_back("initial_gpu_memory_in_mb");
+  envs.push_back("reallocate_gpu_memory_in_mb");
   envs.push_back("allocator_strategy");
 #elif __clang__
   envs.push_back("use_mkldnn");

python/paddle/dataset/flowers.py

@@ -46,9 +46,9 @@ import six
 from six.moves import cPickle as pickle
 __all__ = ['train', 'test', 'valid']
 
-DATA_URL = 'http://paddlemodels.cdn.bcebos.com/flowers/102flowers.tgz'
-LABEL_URL = 'http://paddlemodels.cdn.bcebos.com/flowers/imagelabels.mat'
-SETID_URL = 'http://paddlemodels.cdn.bcebos.com/flowers/setid.mat'
+DATA_URL = 'http://paddlemodels.bj.bcebos.com/flowers/102flowers.tgz'
+LABEL_URL = 'http://paddlemodels.bj.bcebos.com/flowers/imagelabels.mat'
+SETID_URL = 'http://paddlemodels.bj.bcebos.com/flowers/setid.mat'
 DATA_MD5 = '52808999861908f626f3c1f4e79d11fa'
 LABEL_MD5 = 'e0620be6f572b9609742df49c70aed4d'
 SETID_MD5 = 'a5357ecc9cb78c4bef273ce3793fc85c'

python/paddle/dataset/wmt14.py

@@ -15,7 +15,7 @@
 WMT14 dataset.
 The original WMT14 dataset is too large and a small set of data for set is
 provided. This module will download dataset from
-http://paddlepaddle.cdn.bcebos.com/demo/wmt_shrinked_data/wmt14.tgz and
+http://paddlepaddle.bj.bcebos.com/demo/wmt_shrinked_data/wmt14.tgz and
 parse training set and test set into paddle reader creators.
 
 """

python/paddle/fluid/__init__.py

@@ -163,7 +163,8 @@ def __bootstrap__():
 
     if core.is_compiled_with_cuda():
         read_env_flags += [
-            'fraction_of_gpu_memory_to_use', 'cudnn_deterministic',
+            'fraction_of_gpu_memory_to_use', 'initial_gpu_memory_in_mb',
+            'reallocate_gpu_memory_in_mb', 'cudnn_deterministic',
             'enable_cublas_tensor_op_math', 'conv_workspace_size_limit',
             'cudnn_exhaustive_search', 'memory_optimize_debug', 'selected_gpus',
             'sync_nccl_allreduce', 'limit_of_tmp_allocation',

python/paddle/fluid/compiler.py

@@ -17,9 +17,10 @@ import os
 import six
 import sys
 from .. import compat as cpt
+from . import framework
+from .framework import cuda_places, cpu_places
 
 from . import core
-from . import framework
 
 __all__ = ['CompiledProgram', 'ExecutionStrategy', 'BuildStrategy']
 
@@ -44,21 +45,6 @@ def _is_pserver_mode(main_program):
     return False
 
 
-def get_available_places(use_cuda):
-    if use_cuda:
-        gpus_env = os.getenv("FLAGS_selected_gpus")
-        if gpus_env:
-            gpus = [int(s) for s in gpus_env.split(",")]
-        else:
-            gpus = [i for i in six.moves.range(core.get_cuda_device_count())]
-        places = [core.CUDAPlace(i) for i in gpus]
-    else:
-        cpu_num = int(os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
-        places = [core.CPUPlace() for _ in six.moves.range(cpu_num)]
-    assert places, "no place for execution"
-    return places
-
-
 class CompiledProgram(object):
     """
     Compiles to Graph for execution.
@@ -117,7 +103,8 @@ class CompiledProgram(object):
                            loss_name=None,
                            build_strategy=None,
                            exec_strategy=None,
-                           share_vars_from=None):
+                           share_vars_from=None,
+                           places=None):
         """Configs the program to run in data parallel way.
 
         Args:
@@ -132,10 +119,18 @@ class CompiledProgram(object):
                 threads are used, how many iterations to clean up the temp
                 variables. For more information, please refer
                 to fluid.ExecutionStrategy. Default None.
-            share_vars_from(CompiledProgram): If provide, this CompiledProgram
+            share_vars_from(CompiledProgram): If provided, this CompiledProgram
                 will share variables from `share_vars_from`. `share_vars_from`
                 must be run by the executor before this CompiledProgram so that
                 vars are ready.
+            places(list(CUDAPlace)|list(CPUPlace)|None): If provided, only compile
+                program in the given places. Otherwise, the places used when compiled 
+                is determined by the Executor, and the places used are controlled 
+                by environment variables: FLAGS_selected_gpus or CUDA_VISIBLE_DEVICES
+                if using GPU; or CPU_NUM if using CPU. For example, if you want to 
+                run on GPU 0 and 1, set places=[fluid.CUDAPlace(0), fluid.CUDAPlace(1)].
+                If you want to run on 2 CPU cores, set places=[fluid.CPUPlace()]*2.  
+
         Returns:
             self
         """
@@ -150,6 +145,12 @@ class CompiledProgram(object):
             self._exec_strategy = ExecutionStrategy()
         if self._build_strategy is None:
             self._build_strategy = BuildStrategy()
+        if places is not None:
+            if not isinstance(places, (list, tuple)):
+                places = [places]
+            self._places = places
+        else:
+            self._places = None
         self._build_strategy.is_distribution = _is_pserver_mode(self._program)
         return self
 
@@ -192,7 +193,15 @@ class CompiledProgram(object):
             self._local_scopes = []
 
         self._exec_strategy.use_cuda = use_cuda
-        self._places = get_available_places(self._exec_strategy.use_cuda)
+        has_set_place = (self._places is not None)
+        if has_set_place:
+            for p in self._places:
+                assert p._type() == self._place._type(), \
+                    "Place type not match. You may set the wrong type of places"
+        else:
+            self._places = cuda_places(
+            ) if self._exec_strategy.use_cuda else cpu_places()
+        assert self._places, "no place for execution"
 
         if self._exec_strategy.num_threads == 0:
             if self._exec_strategy.use_cuda:
@@ -200,9 +209,7 @@ class CompiledProgram(object):
                 # performance. Worth tunning for other models in the future.
                 self._exec_strategy.num_threads = len(self._places) * 4
             else:
-                cpu_num = int(
-                    os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
-                self._exec_strategy.num_threads = cpu_num * 2
+                self._exec_strategy.num_threads = len(self._places) * 2
 
         # FIXME(dzhwinter): enable_inplace should be after memory_optimize
         # if turn on python memory optimize, turn off the inplace_pass.

python/paddle/fluid/contrib/int8_inference/README.md

@@ -45,28 +45,41 @@ You can load INT8 model by load_inference_model [API](https://github.com/PaddleP
 ```
 
 ## 3. Result
-We provide the results of accuracy measurd on [Intel® Xeon® Platinum Gold Processor](https://ark.intel.com/products/120489/Intel-Xeon-Gold-6148-Processor-27-5M-Cache-2-40-GHz- "Intel® Xeon® Gold 6148 Processor") (also known as Intel® Xeon® Skylake6148).
+We provide the results of accuracy and performance measured on Intel(R) Xeon(R) Gold 6271 (single core).
+
+**I. Top-1 Accuracy on Intel(R) Xeon(R) Gold 6271**
 
 | Model  | Dataset  | FP32 Accuracy  | INT8 Accuracy  | Accuracy Diff  |
-| ------------ | ------------ | ------------ | ------------ | ------------ |
-| ResNet-50  | Small  | 72.00%  | 72.00%  |  0.00% |
-| MobileNet-V1  | Small  | 62.00%  | 62.00%  | 0.00%  |
-| ResNet-50  | Full ImageNet Val  |  76.63%  | 76.17%  | 0.46% |
-| MobileNet-V1 | Full ImageNet Val  | 70.78%  | 70.49%  | 0.29%  |
+| :------------: | :------------: | :------------: | :------------: | :------------: |
+| ResNet-50  | Full ImageNet Val  |  76.63%  | 76.23%  | 0.40% |
+| MobileNet-V1 | Full ImageNet Val  | 70.78%  | 70.47%  | 0.31%  |
+
+**II. Throughput on Intel(R) Xeon(R) Gold 6271 (batch size 1 on single core)**
+
+| Model  | Dataset  | FP32 Throughput  | INT8 Throughput  |  Ratio(INT8/FP32)  |
+| :------------: | :------------: | :------------: | :------------: | :------------: |
+| ResNet-50  | Full ImageNet Val  |  11.54 images/s | 32.2 images/s | 2.79 |
+| MobileNet-V1 | Full ImageNet Val  | 49.21 images/s | 108.37 images/s | 2.2  |
 
-Please note that [Small](http://paddle-inference-dist.cdn.bcebos.com/int8/calibration_test_data.tar.gz "Small") is a subset of [full ImageNet validation dataset](http://www.image-net.org/challenges/LSVRC/2012/nnoupb/ILSVRC2012_img_val.tar "full ImageNet validation dataset"). 
+Please note that [full ImageNet validation dataset](http://www.image-net.org/challenges/LSVRC/2012/nnoupb/ILSVRC2012_img_val.tar "full ImageNet validation dataset") can be downloaded by script `test_calibration.py` with `DATASET=full`. 
 
 Notes:
 * The accuracy measurement requires the model with `label`.
-* The INT8 theoretical speedup is ~1.33X on Intel® Xeon® Skylake Server (please refer to `This allows for 4x more input at the cost of 3x more instructions or 33.33% more compute` in  [Reference](https://software.intel.com/en-us/articles/lower-numerical-precision-deep-learning-inference-and-training "Reference")).
+* The INT8 theoretical speedup is 4X on Intel® Xeon® Cascadelake Server (please refer to `providing a theoretical peak compute gain of 4x int8 OPS over fp32 OPS` in  [Reference](https://software.intel.com/en-us/articles/lower-numerical-precision-deep-learning-inference-and-training "Reference")). However, the actual test results at the model level will be less than 4X, and in general the average is about 2X. In addition, the calculation library optimization of batch size 1 is not as good as the large batch size.
 
 ## 4. How to reproduce the results
-* Small dataset
+* Small dataset (Single core)
 ```bash
 FLAGS_use_mkldnn=true python python/paddle/fluid/contrib/tests/test_calibration.py
 ```
 
-* Full dataset
+* Full dataset (Single core)
 ```bash
 FLAGS_use_mkldnn=true DATASET=full python python/paddle/fluid/contrib/tests/test_calibration.py
 ```
+
+* Full dataset (Multi-core)
+```bash
+FLAGS_use_mkldnn=true OMP_NUM_THREADS=20 DATASET=full python python/paddle/fluid/contrib/tests/test_calibration.py
+```
+> Notes: This is an example command with 20 cores by using set `OMP_NUM_THREADS` value.

python/paddle/fluid/contrib/slim/__init__.py

@@ -13,13 +13,4 @@
 # limitations under the License.
 
 from .core import *
-from .graph import *
-from .prune import *
-__all__ = [
-    'build_compressor',
-    'CompressPass',
-    'ImitationGraph',
-    'SensitivePruneStrategy',
-    'MagnitudePruner',
-    'RatioPruner',
-]
+__all__ = ['Compressor', ]

python/paddle/fluid/contrib/slim/core/__init__.py

@@ -14,11 +14,9 @@
 
 from . import config
 from .config import *
-from . import compress_pass
-from .compress_pass import *
+from . import compressor
+from .compressor import *
 from . import strategy
 from .strategy import *
-from . import pass_builder
-from .pass_builder import *
 
-__all__ = config.__all__ + compress_pass.__all__ + strategy.__all__ + pass_builder.__all__
+__all__ = config.__all__ + compressor.__all__ + strategy.__all__

python/paddle/fluid/contrib/slim/core/compress_pass.py

-#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from ....core import CPUPlace
-from ..graph import get_executor
-
-__all__ = ['Context', 'CompressPass']
-
-
-class Context(object):
-    """
-    The context in the process of compression.
-    Args:
-        exe: The executor used to execute graph.
-        graph: The graph to be compressed.
-        scope: The scope used to execute graph.
-        program_exe: The program_exe is used to execute the program
-                     created for modifying the variables in scope.
-    """
-
-    def __init__(self, exe, graph, scope, program_exe=None):
-        # The total number of epoches to be trained.
-        self.epoch = 0
-        # Current epoch
-        self.epoch_id = 0
-        # Current batch
-        self.batch_id = 0
-        self.exe = exe
-        self.graph = graph
-        self.scope = scope
-        self.program_exe = program_exe
-
-
-class CompressPass(object):
-    """
-    The pass used to compress model.
-    Args:
-        place: The device used in compression.
-        data_reader: The data_reader used to run graph.
-        data_feeder: The data_feeder used to run graph.
-        scope: The scope used to run graph.
-        metrics: The metrics for evaluating model.
-        epoch: The total epoches of trainning in compression.
-        program_exe: The program_exe is used to execute the program
-                     created for modifying the variables in scope.
-    """
-
-    def __init__(self,
-                 place=None,
-                 data_reader=None,
-                 data_feeder=None,
-                 scope=None,
-                 metrics=None,
-                 epoch=None,
-                 program_exe=None):
-        self.strategies = []
-        self.place = CPUPlace() if place is None else place
-        self.data_reader = data_reader
-        self.data_feeder = data_feeder
-        self.scope = scope
-        self.metrics = metrics
-        self.epoch = epoch
-        self.program_exe = program_exe
-
-    def add_strategy(self, strategy):
-        """
-        Add a strategy to current compress pass.
-        Args:
-            strategy: The strategy to be added into current compress pass.
-        """
-        self.strategies.append(strategy)
-        self.epoch = max(strategy.end_epoch, self.epoch)
-
-    def apply(self, graph):
-        """
-        Compress a model.
-        Args:
-            graph: The target graph to be compressed.
-        """
-        self.executor = get_executor(graph, self.place)
-        context = Context(
-            self.executor, graph, self.scope, program_exe=self.program_exe)
-
-        for strategy in self.strategies:
-            strategy.on_compress_begin(context)
-
-        for epoch in range(self.epoch):
-
-            for strategy in self.strategies:
-                strategy.on_epoch_begin(context)
-
-            for data in self.data_reader():
-
-                for strategy in self.strategies:
-                    strategy.on_batch_begin(context)
-                fetches = None
-                if self.metrics:
-                    fetches = self.metrics.values()
-                feed = None
-                if self.data_feeder:
-                    feed = self.data_feeder.feed(data)
-                results = self.executor.run(graph,
-                                            fetches=fetches,
-                                            scope=self.scope,
-                                            feed=feed)
-                if results:
-                    print("results: {}".format(
-                        zip(self.metrics.keys(), results)))
-                for strategy in self.strategies:
-                    strategy.on_batch_end(context)
-                context.batch_id += 1
-
-            for strategy in self.strategies:
-                strategy.on_epoch_end(context)
-            context.epoch_id += 1
-
-        for strategy in self.strategies:
-            strategy.on_compress_end(context)

python/paddle/fluid/contrib/slim/core/config.py

@@ -17,8 +17,9 @@ import funcsigs
 import yaml
 from collections import OrderedDict
 from ..prune import *
-from .compress_pass import *
+from ..quantization import *
 from .strategy import *
+from ..distillation import *
 
 __all__ = ['ConfigFactory']
 """This factory is used to create instances by loading and parsing configure file with yaml format.
@@ -29,15 +30,10 @@ class ConfigFactory(object):
     def __init__(self, config):
         """Init a factory from configure file."""
         self.instances = {}
+        self.compressor = {}
         self.version = None
         self._parse_config(config)
 
-    def get_compress_pass(self):
-        """
-        Get compress pass from factory.
-        """
-        return self.instance('compress_pass')
-
     def instance(self, name):
         """
         Get instance from factory.
@@ -59,8 +55,16 @@ class ConfigFactory(object):
             args = {}
             for key in keys:
                 value = attrs[key]
+                if isinstance(value, str) and value.lower() == 'none':
+                    value = None
                 if isinstance(value, str) and value in self.instances:
                     value = self.instances[value]
+                if isinstance(value, list):
+                    for i in range(len(value)):
+                        if isinstance(value[i],
+                                      str) and value[i] in self.instances:
+                            value[i] = self.instances[value[i]]
+
                 args[key] = value
             self.instances[name] = class_(**args)
         return self.instances.get(name)
@@ -76,16 +80,23 @@ class ConfigFactory(object):
                     assert self.version == int(key_values['version'])
 
                 # parse pruners
-                if key == 'pruners' or key == 'strategies':
+                if key == 'distillers' or key == 'pruners' or key == 'quantizers' or key == 'strategies':
                     instances = key_values[key]
                     for name in instances:
                         self._new_instance(name, instances[name])
 
-                if key == 'compress_pass':
-                    compress_pass = self._new_instance(key, key_values[key])
-                    for name in key_values[key]['strategies']:
-                        strategy = self.instance(name)
-                        compress_pass.add_strategy(strategy)
+                if key == 'compressor':
+                    self.compressor['strategies'] = []
+                    self.compressor['epoch'] = key_values[key]['epoch']
+                    if 'init_model' in key_values[key]:
+                        self.compressor['init_model'] = key_values[key][
+                            'init_model']
+                    self.compressor['checkpoint_path'] = key_values[key][
+                        'checkpoint_path']
+                    if 'strategies' in key_values[key]:
+                        for name in key_values[key]['strategies']:
+                            strategy = self.instance(name)
+                            self.compressor['strategies'].append(strategy)
 
                 if key == 'include':
                     for config_file in key_values[key]:

python/paddle/fluid/contrib/slim/core/pass_builder.py

-#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-from .compress_pass import CompressPass
-from .config import ConfigFactory
-
-__all__ = ['build_compressor']
-
-
-def build_compressor(place=None,
-                     data_reader=None,
-                     data_feeder=None,
-                     scope=None,
-                     metrics=None,
-                     epoch=None,
-                     config=None):
-    if config is not None:
-        factory = ConfigFactory(config)
-        comp_pass = factory.get_compress_pass()
-    else:
-        comp_pass = CompressPass()
-    comp_pass.place = place
-    comp_pass.data_reader = data_reader
-    comp_pass.data_feeder = data_feeder
-    comp_pass.scope = scope
-    comp_pass.metrics = metrics
-    comp_pass.epoch = epoch
-    return comp_pass

python/paddle/fluid/contrib/slim/core/strategy.py

@@ -20,7 +20,7 @@ class Strategy(object):
     Base class for all strategies.
     """
 
-    def __init__(self, start_epoch=0, end_epoch=10):
+    def __init__(self, start_epoch=0, end_epoch=0):
         """
         Args:
             start_epoch: The first epoch to apply the strategy.
@@ -29,7 +29,7 @@ class Strategy(object):
         self.start_epoch = start_epoch
         self.end_epoch = end_epoch
 
-    def on_compress_begin(self, context):
+    def on_compression_begin(self, context):
         pass
 
     def on_epoch_begin(self, context):
@@ -44,5 +44,5 @@ class Strategy(object):
     def on_batch_end(self, context):
         pass
 
-    def on_compress_end(self, context):
+    def on_compression_end(self, context):
         pass

python/paddle/fluid/contrib/slim/demo/filter_prune/config.yaml

-version: 1.0
-pruners:
-    pruner_1:
-        class: 'RatioPruner'
-        ratios:
-            'conv1_1.w': 0.3
-            'conv1_2.w': 0.4
-            '*': 0.9
-        group_dims:
-            '*': [1, 2, 3]
-        criterions:
-            '*': 'l1-norm'
-strategies:
-    strategy_1:
-        class: 'SensitivePruneStrategy'
-        pruner: 'pruner_1'
-        start_epoch: 0
-        end_epoch: 10
-        delta_rate: 0.20
-        acc_loss_threshold: 0.2
-        sensitivities:
-            'conv1_1.w': 0.4
-
-compress_pass:
-    class: 'CompressPass'
-    epoch: 100
-    strategies:
-        - strategy_1

python/paddle/fluid/contrib/slim/demo/filter_prune/demo.py

-# Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import paddle.fluid as fluid
-import paddle
-import os
-import sys
-from paddle.fluid.contrib.slim import CompressPass
-from paddle.fluid.contrib.slim import build_compressor
-from paddle.fluid.contrib.slim import ImitationGraph
-
-
-class LinearModel(object):
-    def __init__(slef):
-        pass
-
-    def train(self):
-        train_program = fluid.Program()
-        startup_program = fluid.Program()
-        startup_program.random_seed = 10
-        with fluid.program_guard(train_program, startup_program):
-            x = fluid.layers.data(name='x', shape=[13], dtype='float32')
-            y = fluid.layers.data(name='y', shape=[1], dtype='float32')
-            predict = fluid.layers.fc(input=x, size=1, act=None)
-            cost = fluid.layers.square_error_cost(input=predict, label=y)
-            avg_cost = fluid.layers.mean(cost)
-            eval_program = train_program.clone()
-            sgd_optimizer = fluid.optimizer.SGD(learning_rate=0.001)
-            sgd_optimizer.minimize(avg_cost)
-
-        train_reader = paddle.batch(
-            paddle.dataset.uci_housing.train(), batch_size=1)
-        eval_reader = paddle.batch(
-            paddle.dataset.uci_housing.test(), batch_size=1)
-        place = fluid.CPUPlace()
-        train_feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
-        eval_feeder = fluid.DataFeeder(place=place, feed_list=[x, y])
-        exe = fluid.Executor(place)
-        exe.run(startup_program)
-        train_metrics = {"loss": avg_cost.name}
-        eval_metrics = {"loss": avg_cost.name}
-
-        graph = ImitationGraph(train_program)
-        config = './config.yaml'
-        comp_pass = build_compressor(
-            place,
-            data_reader=train_reader,
-            data_feeder=train_feeder,
-            scope=fluid.global_scope(),
-            metrics=train_metrics,
-            epoch=1,
-            config=config)
-        comp_pass.apply(graph)
-
-
-if __name__ == "__main__":
-    model = LinearModel()
-    model.train()

python/paddle/fluid/contrib/slim/graph/graph_pass.py → python/paddle/fluid/contrib/slim/distillation/__init__.py

-#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,31 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__all__ = ['GraphPass', 'PruneParameterPass']
+from . import distiller
+from .distiller import *
+from . import distillation_strategy
+from .distillation_strategy import *
 
-
-class GraphPass(object):
-    """
-    Base class for all graph pass.
-    """
-
-    def __init__(self):
-        pass
-
-    def apply(self, graph):
-        pass
-
-
-class PruneParameterPass(GraphPass):
-    """
-    Generate a graph for pruning parameters from target graph.
-    """
-
-    def __init__(self, pruned_params, thresholds):
-        super(PruneParameterPass, self).__init__()
-        self.pruned_params = pruned_params
-        self.thresholds = thresholds
-        self.default_threshold = thresholds['*']
-
-    def apply(self, graph):
-        pass
+__all__ = distiller.__all__
+__all__ += distillation_strategy.__all__

python/paddle/fluid/contrib/slim/distillation/distillation_strategy.py

+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ..core.strategy import Strategy
+from ....framework import Program, program_guard
+from .... import Executor
+import logging
+
+__all__ = ['DistillationStrategy']
+
+logging.basicConfig(format='%(asctime)s-%(levelname)s: %(message)s')
+_logger = logging.getLogger(__name__)
+_logger.setLevel(logging.INFO)
+
+
+class DistillationStrategy(Strategy):
+    def __init__(self, distillers=None, start_epoch=0, end_epoch=0):
+        """
+        Args:
+            distillers(list): A list of distiller used to combine student graph and teacher graph
+                              by adding some loss.
+            start_epoch(int): The epoch when to merge student graph and teacher graph for
+                              distillation training. default: 0
+            end_epoch(int): The epoch when to finish distillation training. default: 0
+            
+        """
+        super(DistillationStrategy, self).__init__(start_epoch, end_epoch)
+        self.distillers = distillers
+
+    def on_compression_begin(self, context):
+        # load from checkpoint
+        if context.epoch_id > 0:
+            if context.epoch_id > self.start_epoch and context.epoch_id < self.end_epoch:
+                _logger.info('Restore DistillationStrategy')
+                self._create_distillation_graph(context)
+                _logger.info('Restore DistillationStrategy finish.')
+
+    def on_epoch_begin(self, context):
+        if self.start_epoch == context.epoch_id:
+            _logger.info('DistillationStrategy::on_epoch_begin.')
+            self._create_distillation_graph(context)
+            _logger.info('DistillationStrategy set optimize_graph.')
+
+    def _create_distillation_graph(self, context):
+        """
+        step 1: Merge student graph and teacher graph into distillation graph.
+        step 2: Add loss into distillation graph by distillers.
+        step 3: Append backward ops and optimize ops into distillation graph for training.
+        """
+        # step 1
+        teacher = context.teacher_graphs[0]
+        for var in teacher.program.list_vars():
+            var.stop_gradient = True
+        graph = context.train_graph.clone()
+        graph.merge(teacher)
+        graph.out_nodes['student_loss'] = graph.out_nodes['loss']
+
+        # step 2
+        for distiller in self.distillers:
+            graph = distiller.distiller_loss(graph)
+
+        # step 3
+        startup_program = Program()
+        with program_guard(graph.program, startup_program):
+            context.distiller_optimizer._name = 'distillation_optimizer'
+            context.distiller_optimizer.minimize(
+                graph.var(graph.out_nodes['loss'])._var)
+        exe = Executor(context.place)
+        exe.run(startup_program, scope=context.scope)
+
+        # backup graph for fine-tune after distillation
+        context.put('distillation_backup_optimize_graph',
+                    context.optimize_graph)
+        context.optimize_graph = graph
+
+    def on_epoch_end(self, context):
+        if context.epoch_id == (self.end_epoch - 1):
+            _logger.info('DistillationStrategy::on_epoch_end.')
+            # restore optimize_graph for fine-tune or other strategy in next stage.
+            context.optimize_graph = context.get(
+                'distillation_backup_optimize_graph')
+            _logger.info(
+                'DistillationStrategy set context.optimize_graph to None.')

python/paddle/fluid/contrib/slim/distillation/distiller.py

+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .... import layers
+from .... import optimizer
+from .... import Executor
+from .... import Program
+from .... import program_guard
+from .... import regularizer
+
+__all__ = ['FSPDistiller', 'L2Distiller']
+
+
+class L2Distiller(object):
+    """
+    Combine two layers from student net and teacher net by l2-loss.
+    And add the loss into the total loss using for distillation training.
+    """
+
+    def __init__(self,
+                 student_feature_map,
+                 teacher_feature_map,
+                 distillation_loss_weight=1):
+        """
+        Args:
+            student_feature_map(str): The name of feature map from student network.
+            teacher_feature_map(str): The name of feature map from teacher network.
+                                      It's shape should be the same with student network.
+            distillation_loss_weight(float): The weight of the l2-loss.
+        """
+        self.student_feature_map = student_feature_map
+        self.teacher_feature_map = teacher_feature_map
+        self.distillation_loss_weight = distillation_loss_weight
+
+    def distiller_loss(self, graph):
+        """
+        Modify graph inplace to add l2-loss.
+        Args: 
+            graph(GraphWrapper): The graph to be modified.
+        Returns:
+            GraphWrapper: The modified graph.
+        """
+        distiller_pass = L2DistillerPass(self.student_feature_map,
+                                         self.teacher_feature_map,
+                                         self.distillation_loss_weight)
+        dis_graph = distiller_pass.apply(graph)
+        return dis_graph
+
+
+class L2DistillerPass(object):
+    """
+    The pass used to add l2-loss.
+    """
+
+    def __init__(self,
+                 student_feature_map,
+                 teacher_feature_map,
+                 distillation_loss_weight=1):
+        """
+        Args:
+            student_feature_map(str): The name of feature map from student network.
+            teacher_feature_map(str): The name of feature map from teacher network.
+                                      It's shape should be the same with student network.
+            distillation_loss_weight(float): The weight of the l2-loss.
+        """
+        self.student_feature_map = student_feature_map
+        self.teacher_feature_map = teacher_feature_map
+        self.distillation_loss_weight = distillation_loss_weight
+
+    def apply(self, graph):
+        ret_graph = graph
+        with program_guard(ret_graph.program):
+
+            student_feature_map = ret_graph.var(self.student_feature_map)._var
+            teacher_feature_map = ret_graph.var(self.teacher_feature_map)._var
+            l2loss = layers.reduce_mean(
+                layers.square(student_feature_map - teacher_feature_map))
+
+            distillation_loss = l2loss * self.distillation_loss_weight
+            student_loss = ret_graph.var(ret_graph.out_nodes['loss'])._var
+            loss = distillation_loss + student_loss
+
+            ret_graph.out_nodes[
+                'l2loss_' + self.student_feature_map + "_" +
+                self.teacher_feature_map] = distillation_loss.name
+            ret_graph.out_nodes['loss'] = loss.name
+        return ret_graph
+
+
+class FSPDistiller(object):
+    """
+    Combine layers from student net and teacher net by fsp-loss.
+    """
+
+    def __init__(self, student_pairs, teacher_pairs,
+                 distillation_loss_weight=1):
+        """
+        Args:
+            student_pairs(list<tuple>): Each tuple, with two variable names, in student_pairs indicates
+                                        a section in student network. The variables in a tuple should
+                                        have the same feature map size.
+            teacher_pairs(list<tuple>): Each tuple, with two variable names, in teacher_pairs indicates
+                                        a section in teacher network. The variables in a tuple should
+                                        have the same feature map size. Varibale named teacher_pairs[i][j]
+                                        should has the save channel number with that of variable named 
+                                        student_pairs[i][j].
+
+            distillation_loss_weight(float): The weight of the fsp-loss. default: 1.
+        """
+        self.student_pairs = student_pairs
+        self.teacher_pairs = teacher_pairs
+        self.distillation_loss_weight = distillation_loss_weight
+
+    def distiller_loss(self, graph):
+        """
+        Modify graph inplace to add fsp-loss.
+        Args: 
+            graph(GraphWrapper): The graph to be modified.
+        Returns:
+            GraphWrapper: The modified graph.
+        """
+        distiller_pass = FSPDistillerPass(self.student_pairs,
+                                          self.teacher_pairs,
+                                          self.distillation_loss_weight)
+        dis_graph = distiller_pass.apply(graph)
+        return dis_graph
+
+
+class FSPDistillerPass(object):
+    '''
+    Combine layers from student net and teacher net by fsp-loss.
+    '''
+
+    def __init__(self, s_pairs, t_pairs, distillation_loss_weight=1):
+        """
+        Args:
+            s_pairs(list<tuple>): Each tuple, with two variable names, in student_pairs indicates
+                                        a section in student network. The variables in a tuple should
+                                        have the same feature map size.
+            t_pairs(list<tuple>): Each tuple, with two variable names, in teacher_pairs indicates
+                                        a section in teacher network. The variables in a tuple should
+                                        have the same feature map size. Varibale named teacher_pairs[i][j]
+                                        should has the save channel number with that of variable named 
+                                        student_pairs[i][j].
+
+            distillation_loss_weight(float): The weight of the fsp-loss. default: 1.
+        """
+        self.s_pairs = s_pairs
+        self.t_pairs = t_pairs
+        self.distillation_loss_weight = distillation_loss_weight
+
+    def apply(self, graph):
+        ret_graph = graph
+        with program_guard(ret_graph.program):
+            losses = []
+            for s_pair, t_pair in zip(self.s_pairs, self.t_pairs):
+                s_pair_start = ret_graph.var(s_pair[0])._var
+                s_pair_end = ret_graph.var(s_pair[1])._var
+                s_fsp_matrix = self._fsp_matrix(s_pair_start, s_pair_end)
+                t_pair_start = ret_graph.var(t_pair[0])._var
+                t_pair_end = ret_graph.var(t_pair[1])._var
+                t_fsp_matrix = self._fsp_matrix(t_pair_start, t_pair_end)
+                l2_loss = layers.reduce_mean(
+                    layers.square(s_fsp_matrix - t_fsp_matrix))
+                losses.append(l2_loss)
+            distillation_loss = layers.sum(
+                losses) * self.distillation_loss_weight
+            student_loss = ret_graph.var(ret_graph.out_nodes['loss'])._var
+            loss = distillation_loss + student_loss
+
+            ret_graph.out_nodes[
+                'fsp_distillation_loss'] = distillation_loss.name
+            ret_graph.out_nodes['loss'] = loss.name
+        return ret_graph
+
+    def _fsp_matrix(self, fea_map_0, fea_map_1):
+        return layers.fsp_matrix(fea_map_0, fea_map_1)

python/paddle/fluid/contrib/slim/graph/__init__.py

@@ -14,10 +14,7 @@
 
 from . import executor
 from .executor import *
-from . import graph
-from .graph import *
-from . import graph_pass
-from .graph_pass import *
+from . import graph_wrapper
+from .graph_wrapper import *
 __all__ = executor.__all__
-__all__ += graph.__all__
-__all__ += graph_pass.__all__
+__all__ += graph_wrapper.__all__

python/paddle/fluid/contrib/slim/graph/executor.py

@@ -12,51 +12,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import abc
-from abc import abstractmethod
+from ....compiler import CompiledProgram
+from ....data_feeder import DataFeeder
 from .... import executor
-from .graph import IRGraph, ImitationGraph
+from .graph_wrapper import GraphWrapper
 
-__all__ = ['get_executor']
+__all__ = ['SlimGraphExecutor']
 
 
-class GraphExecutor(object):
-    __metaclass__ = abc.ABCMeta
+class SlimGraphExecutor(object):
+    """
+    Wrapper of executor used to run GraphWrapper.
+    """
 
     def __init__(self, place):
-        self.place = place
-
-    @abstractmethod
-    def run(self, graph, feches=None, feed=None):
-        pass
-
-
-class IRGraphExecutor(GraphExecutor):
-    def run(self, grah, fetches, feed=None):
-        pass
-
-
-class ImitationGraphExecutor(GraphExecutor):
-    def __init__(self, place):
-        super(ImitationGraphExecutor, self).__init__(place)
         self.exe = executor.Executor(place)
+        self.place = place
 
-    def run(self, graph, scope=None, fetches=None, feed=None):
-        assert isinstance(graph, ImitationGraph)
-        fetch_list = None
-        if fetches:
-            fetch_list = [
-                graph.program.global_block().var(name) for name in fetches
-            ]
-        results = self.exe.run(graph.program,
+    def run(self, graph, scope, data=None):
+        """
+        Runing a graph with a batch of data.
+        Args:
+            graph(GraphWrapper): The graph to be executed.
+            scope(fluid.core.Scope): The scope to be used.
+            data(list<tuple>): A batch of data. Each tuple in this list is a sample.
+                               It will feed the items of tuple to the in_nodes of graph.
+        Returns:
+            results(list): A list of result with the same order indicated by graph.out_nodes.
+        """
+        assert isinstance(graph, GraphWrapper)
+        if data is not None:
+            feeder = DataFeeder(
+                feed_list=graph.in_nodes.values(),
+                place=self.place,
+                program=graph.program)
+            feed = feeder.feed(data)
+
+        fetch_list = graph.out_nodes.values()
+        program = graph.compiled_graph if graph.compiled_graph else graph.program
+        results = self.exe.run(program,
                                scope=scope,
                                fetch_list=fetch_list,
                                feed=feed)
         return results
-
-
-def get_executor(graph, place):
-    if isinstance(graph, ImitationGraph):
-        return ImitationGraphExecutor(place)
-    if isinstance(graph, IRGraph):
-        return IRGraphExecutor(place)

python/paddle/fluid/contrib/slim/graph/graph_wrapper.py

+# Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from collections import OrderedDict
+from .... import io
+from .... import compiler
+from ....framework import Program
+from ....framework import program_guard
+from ....framework import Parameter
+from ....framework import Variable
+from ....executor import Executor
+import copy
+from collections import Iterable
+from ....io import save_inference_model, load_inference_model, save_persistables
+import numpy as np
+import pickle
+import os
+
+__all__ = ['GraphWrapper', 'VarWrapper', 'OpWrapper']
+
+OPTIMIZER_OPS = [
+    'momentum',
+    'lars_momentum',
+    'adagrad',
+    'adam',
+    'adamax',
+    'decayed_adagrad',
+    'adadelta',
+    'rmsprop',
+]
+
+
+class VarWrapper(object):
+    def __init__(self, var, graph):
+        assert isinstance(var, Variable)
+        assert isinstance(graph, GraphWrapper)
+        self._var = var
+        self._graph = graph
+
+    def __eq__(self, v):
+        """
+        Overwrite this function for ...in... syntax in python.
+        """
+        return self._var.name == v._var.name
+
+    def name(self):
+        """
+        Get the name of the variable.
+        """
+        return self._var.name
+
+    def shape(self):
+        """
+        Get the shape of the varibale.
+        """
+        return self._var.shape
+
+    def set_shape(self, shape):
+        """
+        Set the shape of the variable.
+        """
+        self._var.desc.set_shape(shape)
+
+    def inputs(self):
+        """
+        Get all the operators that use this variable as output.
+        Returns:
+            list<OpWrapper>: A list of operators.
+        """
+        ops = []
+        for op in self._graph.ops():
+            if self in op.all_inputs():
+                ops.append(op)
+        return ops
+
+    def outputs(self):
+        """
+        Get all the operators that use this variable as input.
+        Returns:
+            list<OpWrapper>: A list of operators.
+        """
+        ops = []
+        for op in self._graph.ops():
+            if self in op.all_outputs():
+                ops.append(op)
+        return ops
+
+
+class OpWrapper(object):
+    def __init__(self, op, graph):
+        assert isinstance(graph, GraphWrapper)
+        self._op = op
+        self._graph = graph
+
+    def __eq__(self, op):
+        """
+        Overwrite this function for ...in... syntax in python.
+        """
+        return self.idx() == op.idx()
+
+    def all_inputs(self):
+        """
+        Get all the input variables of this operator.
+        """
+        return [
+            self._graph.var(var_name) for var_name in self._op.input_arg_names
+        ]
+
+    def all_outputs(self):
+        """
+        Get all the output variables of this operator.
+        """
+        return [
+            self._graph.var(var_name) for var_name in self._op.output_arg_names
+        ]
+
+    def idx(self):
+        """
+        Get the id of this operator.
+        """
+        return self._op.idx
+
+    def type(self):
+        """
+        Get the type of this operator.
+        """
+        return self._op.type
+
+    def is_bwd_op(self):
+        """
+        Whether this operator is backward op.
+        """
+        return self.type().endswith('_grad')
+
+    def is_opt_op(self):
+        """
+        Whether this operator is optimizer op.
+        """
+        return self.type() in OPTIMIZER_OPS
+
+    def inputs(self, name):
+        """
+        Get all the varibales by the input name.
+        """
+        return [self._graph.var(var_name) for var_name in self._op.input(name)]
+
+    def outputs(self, name):
+        """
+        Get all the varibales by the output name.
+        """
+        return [self._graph.var(var_name) for var_name in self._op.output(name)]
+
+    def set_attr(self, key, value):
+        """
+        Set the value of attribute by attribute's name.
+
+        Args:
+            key(str): the attribute name.
+            value(bool|int|str|float|list): the value of the attribute.
+        """
+        self._op._set_attr(key, value)
+
+    def attr(self, name):
+        """
+        Get the attribute by name.
+
+        Args:
+            name(str): the attribute name.
+
+        Returns:
+            bool|int|str|float|list: The attribute value. The return value
+            can be any valid attribute type.
+        """
+        return self._op.attr(name)
+
+
+class GraphWrapper(object):
+    """
+    It is a wrapper of paddle.fluid.framework.IrGraph with some special functions
+    for paddle slim framework.
+    """
+
+    def __init__(self, program=None, in_nodes=[], out_nodes=[]):
+        """
+        Args:
+            program(framework.Program): A program with 
+            in_nodes(dict): A dict to indicate the input nodes of the graph.
+                            The key is user-defined and human-readable name.
+                            The value is the name of Variable.
+            out_nodes(dict): A dict to indicate the input nodes of the graph.
+                            The key is user-defined and human-readable name.
+                            The value is the name of Variable.
+        """
+        super(GraphWrapper, self).__init__()
+        self.program = Program() if program is None else program
+        self.compiled_graph = None
+        self.in_nodes = OrderedDict(in_nodes)
+        self.out_nodes = OrderedDict(out_nodes)
+        self._attrs = OrderedDict()
+
+    def all_parameters(self):
+        """
+        Get all the parameters in this graph.
+        Returns:
+            list<VarWrapper>: A list of VarWrapper instances.
+        """
+        params = []
+        for block in self.program.blocks:
+            for param in block.all_parameters():
+                params.append(VarWrapper(param, self))
+        return params
+
+    def is_parameter(self, var):
+        """
+        Whether the given variable is parameter.
+        Args:
+            var(VarWrapper): The given varibale.
+        """
+        return isinstance(var._var, Parameter)
+
+    def is_persistable(self, var):
+        """
+        Whether the given variable is persistable.
+        Args:
+            var(VarWrapper): The given varibale.
+        """
+        return var._var.persistable
+
+    def compile(self, for_parallel=True, for_test=False):
+        """
+        Compile the program in this wrapper to framework.CompiledProgram for next running.
+        This function must be called if the program is modified.
+        Args:
+            for_parallel(bool): Whether the program to run in data parallel way. default: True.
+            for_test(bool): Whether the compiled program is used for test.
+        """
+        target = self.program
+        if for_test:
+            loss = None
+        else:
+            loss = self.out_nodes['loss']
+        if for_parallel:
+            # disable memory optimize for stable training
+            build_strategy = compiler.BuildStrategy()
+            build_strategy.enable_inplace = False
+            build_strategy.memory_optimize = False
+            self.compiled_graph = compiler.CompiledProgram(
+                target).with_data_parallel(
+                    loss_name=loss, build_strategy=build_strategy)
+        else:
+            self.compiled_graph = compiler.CompiledProgram(target)
+
+    def ops(self):
+        """
+        Return all operator nodes included in the graph as a set.
+        """
+        ops = []
+        for block in self.program.blocks:
+            for op in block.ops:
+                ops.append(OpWrapper(op, self))
+        return ops
+
+    def vars(self):
+        """
+        Get all the variables.
+        """
+        return [VarWrapper(var, self) for var in self.program.list_vars()]
+
+    def var(self, name):
+        """
+        Get the variable by variable name.
+        """
+        return VarWrapper(self.program.global_block().var(name), self)
+
+    def clone(self, for_test=False):
+        """
+        Clone a new graph from current graph.
+        Returns:
+            (GraphWrapper): The wrapper of a new graph.
+        """
+        return GraphWrapper(
+            self.program.clone(for_test),
+            copy.deepcopy(self.in_nodes), copy.deepcopy(self.out_nodes))
+
+    def merge(self, graph):
+        """
+        Merge a graph into current graph.
+        Args:
+            graph(GraphWrapper): The graph to be merged by current graph.
+        """
+        for var in graph.program.list_vars():
+            new_var = self.program.global_block()._clone_variable(
+                var, force_persistable=False)
+            new_var.stop_gradient = var.stop_gradient
+            # TODO: parameters should be cloned
+        for op in graph.ops():
+            op = op._op
+            inputs = {}
+            outputs = {}
+            attrs = {}
+            for input_name in op.input_names:
+                inputs[input_name] = [
+                    self.var(in_var_name)._var
+                    for in_var_name in op.input(input_name)
+                ]
+            for output_name in op.output_names:
+                outputs[output_name] = [
+                    self.var(out_var_name)._var
+                    for out_var_name in op.output(output_name)
+                ]
+            for attr_name in op.attr_names:
+                attrs[attr_name] = op.attr(attr_name)
+            self.program.global_block().append_op(
+                type=op.type, inputs=inputs, outputs=outputs, attrs=attrs)
+
+    def program(self):
+        """
+        Get the program in current wrapper.
+        """
+        return self.program
+
+    def pre_ops(self, op):
+        """
+        Get all the previous operators of target operator.
+        Args:
+            op(OpWrapper): Target operator..
+        Returns:
+            list<OpWrapper>: A list of operators.
+        """
+        ops = []
+        for p in self.ops():
+            for in_var in op.all_inputs():
+                if in_var in p.all_outputs():
+                    ops.append(p)
+        return ops
+
+    def next_ops(self, op):
+        """
+        Get all the next operators of target operator.
+        Args:
+            op(OpWrapper): Target operator..
+        Returns:
+            list<OpWrapper>: A list of operators.
+        """
+        ops = []
+        for p in self.ops():
+            for out_var in op.all_outputs():
+                if out_var in p.all_inputs():
+                    ops.append(p)
+        return ops
+
+    def get_param_by_op(self, op):
+        """
+        Get the parameters used by target operator.
+        """
+        assert isinstance(op, OpWrapper)
+        params = []
+        for var in op.all_inputs():
+            if isinstance(var._var, Parameter):
+                params.append(var)
+        assert len(params) > 0
+        return params
+
+    def numel_params(self):
+        """
+        Get the number of elements in all parameters.
+        """
+        ret = 0
+        for param in self.all_parameters():
+            ret += np.product(param.shape())
+        return ret
+
+    def get_optimize_graph(self, optimizer, place, scope, no_grad_var_names=[]):
+        """
+        Get a new graph for training by appending some backward operators and optimization operators.
+        Args:
+            optimizer: The optimzier used to generate training graph.
+            place: The place to run the graph.
+            scope: The scope used to run the graph. Some new variable will be added into this scope.
+            no_grad_var_names(list<str>): Names of variables that should be ignored while computing gradients. default: [].
+        Returns:
+            (GraphWrapper): The wrapper of new graph with backward ops and optimization ops. 
+        """
+        graph = self.clone()
+        startup_program = Program()
+        with program_guard(
+                main_program=graph.program, startup_program=startup_program):
+            target_name = None
+            if 'loss' in graph.out_nodes:
+                target_name = graph.out_nodes['loss']
+            elif 'cost' in graph.out_nodes:
+                target_name = graph.out_nodes['cost']
+            target = graph.var(target_name)._var
+            optimizer.minimize(target, no_grad_set=no_grad_var_names)
+
+        exe = Executor(place)
+        exe.run(program=startup_program, scope=scope)
+        return graph
+
+    def flops(self, only_conv=False):
+        """
+        Get the flops of current graph.
+        Args:
+            only_conv: Only calculating the conv layers. default: False.
+        Returns:
+            int: The flops of current graph.
+        """
+        flops = 0
+        for op in self.ops():
+            if op.type() in ['conv2d', 'depthwise_conv2d']:
+                filter_shape = op.inputs("Filter")[0].shape()
+                input_shape = op.inputs("Input")[0].shape()
+                output_shape = op.outputs("Output")[0].shape()
+                c_out, c_in, k_h, k_w = filter_shape
+                _, _, h_out, w_out = output_shape
+                groups = op.attr("groups")
+                kernel_ops = k_h * k_w * (c_in / groups)
+                if len(op.inputs("Bias")) > 0:
+                    with_bias = 1
+                else:
+                    with_bias = 0
+                flops += 2 * h_out * w_out * c_out * (kernel_ops + with_bias)
+            elif op.type() == 'pool2d' and not only_conv:
+                input_shape = op.inputs("X")[0].shape()
+                output_shape = op.outputs("Out")[0].shape()
+                _, c_out, h_out, w_out = output_shape
+                k_size = op.attr("ksize")
+                flops += h_out * w_out * c_out * (k_size[0]**2)
+
+            elif op.type() == 'mul' and not only_conv:
+                x_shape = list(op.inputs("X")[0].shape())
+                y_shape = op.inputs("Y")[0].shape()
+                if x_shape[0] == -1:
+                    x_shape[0] = 1
+                flops += 2 * x_shape[0] * x_shape[1] * y_shape[1]
+
+            elif op.type() in ['relu', 'sigmoid', 'batch_norm'
+                               ] and not only_conv:
+                input_shape = list(op.inputs("X")[0].shape())
+                if input_shape[0] == -1:
+                    input_shape[0] = 1
+                flops += np.product(input_shape)
+
+        return flops
+
+    def save_persistables(self, path, exe):
+        """
+        Save all the persistable variables into file.
+        Args:
+            path(str): The path to save the persistables.
+            exe(framework.Executor): The executor used to save the persistables.
+        """
+        io.save_persistables(exe.exe, path, main_program=self.program)
+
+    def load_persistables(self, path, exe):
+        """
+        Load the persistable variables from file.
+        Args:
+            path(str): The path to load the persistables.
+            exe(framework.Executor): The executor used to load the persistables.
+        """
+
+        def if_exist(var):
+            return os.path.exists(os.path.join(path, var.name))
+
+        io.load_vars(
+            exe.exe, path, main_program=self.program, predicate=if_exist)
+
+    def update_param_shape(self, scope):
+        """
+        Update the shape of parameters in the graph according to tensors in scope.
+        It is used after loading pruned parameters from file.
+        """
+        for param in self.all_parameters():
+            tensor_shape = np.array(scope.find_var(param.name()).get_tensor(
+            )).shape
+            param.set_shape(tensor_shape)
+
+    def infer_shape(self):
+        """
+        Update the groups of convolution layer according to current filters.
+        It is used after loading pruned parameters from file.
+        """
+        for op in self.ops():
+            if op.type() != 'conditional_block':
+                op._op.desc.infer_shape(op._op.block.desc)
+
+    def update_groups_of_conv(self):
+        for op in self.ops():
+            if op.type() == 'depthwise_conv2d':
+                op.set_attr('groups', op.inputs('Filter')[0].shape()[0])

python/paddle/fluid/contrib/slim/prune/pruner.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 import numpy as np
+import collections
 from .... import layers
 
-__all__ = ['Pruner', 'MagnitudePruner', 'RatioPruner']
+__all__ = ['Pruner', 'StructurePruner']
 
 
 class Pruner(object):
@@ -30,54 +31,77 @@ class Pruner(object):
         pass
 
 
-class MagnitudePruner(Pruner):
+class StructurePruner(Pruner):
     """
-    Pruner used to pruning a parameter by threshold.
+    Pruner used to pruning parameters by groups.
     """
 
-    def __init__(self, threshold):
-        self.threshold = threshold
-
-    def prune(self, param, threshold=None):
-        if threshold is None:
-            thres = layers.fill_constant(
-                shape=[1], dtype='float32', value=self.threshold)
-        else:
-            thres = threshold
-        zeros_mask = layers.less_than(x=param, y=thres)
-        return zeros_mask
-
-
-class RatioPruner(Pruner):
-    """
-    Pruner used to pruning a parameter by ratio.
-    """
+    def __init__(self, pruning_axis, criterions):
+        """
+        Args:
+            pruning_axis(dict): The key is the name of parameter to be pruned,
+                                '*' means all the parameters.
+                                The value is the axis to be used. Given a parameter
+                                with shape [3, 4], the result of pruning 50% on aixs 1
+                                is a parameter with shape [3, 2].
+            criterions(dict): The key is the name of parameter to be pruned,
+                              '*' means all the parameters.
+                              The value is the criterion used to sort groups for pruning.
+                              It only supports 'l1_norm' currently.
+        """
+        self.pruning_axis = pruning_axis
+        self.criterions = criterions
 
-    def __init__(self, ratios=None):
+    def cal_pruned_idx(self, name, param, ratio, axis=None):
         """
+        Calculate the index to be pruned on axis by given pruning ratio.
         Args:
-            ratios: dict with pair (paramer_name, pruned_ratio). 
+            name(str): The name of parameter to be pruned.
+            param(np.array): The data of parameter to be pruned.
+            ratio(float): The ratio to be pruned.
+            axis(int): The axis to be used for pruning given parameter.
+                       If it is None, the value in self.pruning_axis will be used.
+                       default: None.
+        Returns:
+            list<int>: The indexes to be pruned on axis.
         """
-        self.ratios = ratios
+        criterion = self.criterions[
+            name] if name in self.criterions else self.criterions['*']
+        if axis is None:
+            assert self.pruning_axis is not None, "pruning_axis should set if axis is None."
+            axis = self.pruning_axis[
+                name] if name in self.pruning_axis else self.pruning_axis['*']
+        prune_num = int(round(param.shape[axis] * ratio))
+        reduce_dims = [i for i in range(len(param.shape)) if i != axis]
+        if criterion == 'l1_norm':
+            criterions = np.sum(np.abs(param), axis=tuple(reduce_dims))
+        pruned_idx = criterions.argsort()[:prune_num]
+        return pruned_idx
 
-    def prune(self, param, ratio=None):
+    def prune_tensor(self, tensor, pruned_idx, pruned_axis, lazy=False):
         """
+        Pruning a array by indexes on given axis.
         Args:
-            ratio: `ratio=40%` means pruning (1 - 40%) weights to zero.
+            tensor(numpy.array): The target array to be pruned.
+            pruned_idx(list<int>): The indexes to be pruned.
+            pruned_axis(int): The axis of given array to be pruned on. 
+            lazy(bool): True means setting the pruned elements to zero.
+                        False means remove the pruned elements from memory.
+                        default: False.
+        Returns:
+            numpy.array: The pruned array.
         """
-        if ratio is None:
-            rat = self.ratios[
-                param.name] if param.name in self.ratios else self.ratios['*']
-        else:
-            rat = ratio
-        if rat < 1.0:
-            k = max(int(rat * np.prod(param.shape)), 1)
-            param_vec = layers.reshape(x=param, shape=[1, -1])
-            param_topk, _ = layers.topk(param_vec, k=k)
-            threshold = layers.slice(
-                param_topk, axes=[1], starts=[-1], ends=[k])
-            threshold = layers.reshape(x=threshold, shape=[1])
-            zeros_mask = layers.less_than(x=param, y=threshold)
+        mask = np.zeros(tensor.shape[pruned_axis], dtype=bool)
+        mask[pruned_idx] = True
+
+        def func(data):
+            return data[~mask]
+
+        def lazy_func(data):
+            data[mask] = 0
+            return data
+
+        if lazy:
+            return np.apply_along_axis(lazy_func, pruned_axis, tensor)
         else:
-            zeros_mask = layers.ones(param.shape)
-        return zeros_mask
+            return np.apply_along_axis(func, pruned_axis, tensor)

python/paddle/fluid/contrib/slim/quantization/__init__.py

@@ -16,5 +16,7 @@ from __future__ import print_function
 
 from . import quantization_pass
 from .quantization_pass import *
+from . import quantization_strategy
+from .quantization_strategy import *
 
-__all__ = quantization_pass.__all__
+__all__ = quantization_pass.__all__ + quantization_strategy.__all__

python/paddle/fluid/contrib/slim/tests/configs/config.yaml

-version: 1.0
-include: ["./configs/pruners.yaml", "./configs/pruners_0.yaml"]
-pruners:
-    pruner_1:
-        class: 'RatioPruner'
-        ratios:
-            'conv1_1.w': 0.3
-            'conv1_2.w': 0.4
-            '*': 0.9
-        group_dims:
-            '*': [1, 2, 3]
-        criterions:
-            '*': 'l1-norm'
-strategies:
-    strategy_1:
-        class: 'SensitivePruneStrategy'
-        pruner: 'pruner_2'
-        start_epoch: 0
-        end_epoch: 10
-        delta_rate: 0.20
-        acc_loss_threshold: 0.2
-        sensitivities:
-            'conv1_1.w': 0.4
-
-compress_pass:
-    class: 'CompressPass'
-    epoch: 100
-    strategies:
-        - strategy_1

python/paddle/fluid/contrib/slim/tests/configs/filter_pruning.yaml

+#start_epoch:         The 'on_epoch_begin' function will be called in start_epoch. default: 0.
+#end_epoch:           The 'on_epoch_end' function will be called in end_epoch. default: 10.
+#delta_rate:          The delta used to generate ratios when calculating sensitivities.
+#target_ratio:        The flops ratio to be pruned from current model.
+#metric_name:         The metric used to evaluate the model.
+#pruned_params:       The pattern str to match the parameter names to be pruned.
+#sensitivities_file:  The sensitivities file.
+#num_steps:           The number of pruning steps.
+#eval_rate:           The rate of sampled data used to calculate sensitivities.
+version: 1.0
+pruners:
+    pruner_1:
+        class: 'StructurePruner'
+        pruning_axis:
+            '*': 0
+        criterions:
+            '*': 'l1_norm'
+strategies:
+    sensitive_pruning_strategy:
+        class: 'SensitivePruneStrategy'
+        pruner: 'pruner_1'
+        start_epoch: 0
+        delta_rate: 0.1
+        target_ratio: 0.3
+        num_steps: 1
+        eval_rate: 0.5
+        pruned_params: '.*_sep_weights'
+        sensitivities_file: 'mobilenet_acc_top1_sensitive.data'
+        metric_name: 'acc_top1'
+compressor:
+    epoch: 120
+    checkpoint_path: './checkpoints/'
+    strategies:
+        - sensitive_pruning_strategy

python/paddle/fluid/contrib/slim/tests/configs/pruners.yaml

-version: 1.0
-pruners:
-    pruner_2:
-        class: 'RatioPruner'
-        ratios:
-            'conv1_1.w': 0.5
-            'conv1_2.w': 0.2
-            '*': 0.7
-        group_dims:
-            '*': [1, 2, 3]
-        criterions:
-            '*': 'l1-norm'

python/paddle/fluid/contrib/slim/tests/configs/pruners_0.yaml

-version: 1.0
-pruners:
-    pruner_3:
-        class: 'RatioPruner'
-        ratios:
-            'conv1_1.w': 0.5
-            'conv1_2.w': 0.2
-            '*': 0.7
-        group_dims:
-            '*': [1, 2, 3]
-        criterions:
-            '*': 'l1-norm'

python/paddle/fluid/executor.py

@@ -564,6 +564,10 @@ class Executor(object):
 
         if feed is None:
             feed = {}
+        elif isinstance(feed, (list, tuple)):
+            assert len(feed) == 1, "Not compiled with data parallel"
+            feed = feed[0]
+
         if not isinstance(feed, dict):
             raise TypeError(
                 "feed requires dict as its Parameter. But you passed in %s" %

python/paddle/fluid/imperative/layer_object_helper.py

@@ -105,6 +105,7 @@ class LayerObjectHelper(LayerHelperBase):
 
         Returns dtype of the input
         """
+        inputs_in = inputs_in if (inputs_in is not None) else []
         inputs = self._multiple_input(inputs_in)
         dtype = None
         for each in inputs: