summary | shortlog | log | commit | commitdiff | tree
raw | patch | inline | side by side (parent: 8d5c969)
raw | patch | inline | side by side (parent: 8d5c969)
author | Sergei Nikolaev <snikolaev@nvidia.com> | |
Tue, 9 May 2017 21:57:00 +0000 (14:57 -0700) | ||
committer | Sergei Nikolaev <snikolaev@nvidia.com> | |
Tue, 9 May 2017 22:31:13 +0000 (15:31 -0700) |
index 554adee88c4c56061e68f6a57af2c0549da8e606..05152ab7168d064d3a0a1a2dc517c8ef2b8f7131 100644 (file)
bool debug = false) {
mgr_.init(gpus, m, debug);
}
+ ~Scope() {
+ mgr_.reset();
+ }
};
// Workspace's release() functionality depends on global pool availability
const char* pool_name() const;
bool using_pool() const { return mode_ != CUDA_MALLOC; }
void init(const std::vector<int>&, Mode, bool);
+ void reset();
void* pinned_buffer(size_t size, int device, int group);
std::string report_dev_info(int device);
shared_mutex& read_write_mutex() { return mutex_; }
index b2a9d27f7644544a66291d212019571a2b41b21d..cfb375dacdb24ef76c74d0d05cf5328859f1ca8a 100644 (file)
--- a/python/caffe/__init__.py
+++ b/python/caffe/__init__.py
-from .pycaffe import Net, SGDSolver, NesterovSolver, AdaGradSolver, RMSPropSolver, AdaDeltaSolver, AdamSolver
-from ._caffe import set_mode_cpu, set_mode_gpu, set_device, Layer, get_solver, layer_type_list
+from .pycaffe import Net, SGDSolver, NesterovSolver, AdaGradSolver, RMSPropSolver, \
+ AdaDeltaSolver, AdamSolver
+from ._caffe import set_mode_cpu, set_mode_gpu, set_device, set_devices, Layer, get_solver, \
+ layer_type_list
from ._caffe import CAFFE_VERSION as __version__
from .proto.caffe_pb2 import TRAIN, TEST
from .classifier import Classifier
index 3d5fd60c3ed66a96e1abae78ea2f8c96a8c98295..2bc4e6eaab4c00bebe134ad90593a17c3a316d88 100644 (file)
--- a/python/caffe/_caffe.cpp
+++ b/python/caffe/_caffe.cpp
#ifndef CPU_ONLY
shared_ptr<GPUMemory::Scope> gpu_memory_scope;
+#endif
-void initialize_gpu_memory_scope() {
- vector<int> gpus;
- int count = 0;
- CUDA_CHECK(cudaGetDeviceCount(&count));
- for (int i = 0; i < count; ++i) {
- gpus.push_back(i);
+void initialize_gpu_memory_scope(const vector<int>& gpus) {
+ FLAGS_alsologtostderr = 1;
+ static bool google_initialized = false;
+ if (!google_initialized) {
+ google_initialized = true;
+ google::InitGoogleLogging("pyNVCaffe");
}
- CHECK_GT(gpus.size(), 0);
- Caffe::SetDevice(gpus[0]);
+#ifndef CPU_ONLY
+ gpu_memory_scope.reset();
gpu_memory_scope.reset(new GPUMemory::Scope(gpus));
-}
+ if (gpus.size() > 0) {
+ Caffe::SetDevice(gpus[0]);
+ }
+#else
+ LOG(WARNING) << "GPU memory initialization is ignored for CPU_ONLY build";
#endif
+}
// Selecting mode.
void set_mode_cpu() {
Caffe::set_mode(Caffe::CPU);
#ifndef CPU_ONLY
// We need to run GPU-built Caffe on CPU sometimes.
- initialize_gpu_memory_scope();
+ vector<int> gpus(1, 0);
+ initialize_gpu_memory_scope(gpus);
#endif
}
void set_mode_gpu() {
Caffe::set_mode(Caffe::GPU);
#ifndef CPU_ONLY
- initialize_gpu_memory_scope();
+ vector<int> gpus(1, 0);
+ initialize_gpu_memory_scope(gpus);
+#endif
+}
+
+void set_device(int gpu) {
+ CHECK_GE(gpu, 0);
+ Caffe::set_mode(Caffe::GPU);
+#ifndef CPU_ONLY
+ vector<int> gpus(1, gpu);
+ initialize_gpu_memory_scope(gpus);
#endif
}
+void set_devices(const bp::list& lst) {
+ CHECK(!lst.is_none());
+ boost::python::ssize_t len = boost::python::len(lst);
+ CHECK(len);
+ vector<int> gpus(len);
+ for (int i = 0; i < len; ++i) {
+ gpus[i] = boost::python::extract<int>(lst[i]);
+ CHECK_GE(gpus[i], 0);
+ }
+ Caffe::set_mode(Caffe::GPU);
+ initialize_gpu_memory_scope(gpus);
+}
+
// For convenience, check that input files can be opened, and raise an
// exception that boost will send to Python if not (caffe could still crash
// later if the input files are disturbed before they are actually used, but
// Caffe utility functions
bp::def("set_mode_cpu", &set_mode_cpu);
bp::def("set_mode_gpu", &set_mode_gpu);
- bp::def("set_device", &Caffe::SetDevice);
+ bp::def("set_devices", &set_devices);
+ bp::def("set_device", &set_device);
bp::def("layer_type_list", &LayerRegistry::LayerTypeList);
index 6bf720b138e29ca74b548e6587266b4e29dbdb67..4582407c52a1f82055f6162ecb15ce9a946de0d6 100644 (file)
}
}
+void GPUMemory::Manager::reset() {
+ if (!initialized_) {
+ return;
+ }
+ cub_allocator_.reset();
+ mode_ = CUDA_MALLOC;
+ initialized_ = false;
+ LOG(INFO) << "GPUMemory::Manager uninitialized";
+}
+
GPUMemory::Manager::~Manager() {
for (vector<void*>& buffers_group : pinned_host_buffers_) {
for (void* buffer : buffers_group) {