-
Notifications
You must be signed in to change notification settings - Fork 93
/
pycaffe.cpp
164 lines (149 loc) · 5.93 KB
/
pycaffe.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
// Copyright Yangqing Jia 2013
// pycaffe provides a wrapper of the caffe::Net class as well as some
// caffe::Caffe functions so that one could easily call it from Python.
// Note that for python, we will simply use float as the data type.
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <boost/python.hpp>
#include <numpy/arrayobject.h>
#include "caffe/caffe.hpp"
// Temporary solution for numpy < 1.7 versions: old macro.
#ifndef NPY_ARRAY_C_CONTIGUOUS
#define NPY_ARRAY_C_CONTIGUOUS NPY_C_CONTIGUOUS
#endif
using namespace caffe;
using boost::python::extract;
using boost::python::len;
using boost::python::list;
using boost::python::object;
// A simple wrapper over CaffeNet that runs the forward process.
struct CaffeNet
{
CaffeNet(string param_file, string pretrained_param_file) {
net_.reset(new Net<float>(param_file));
net_->CopyTrainedLayersFrom(pretrained_param_file);
}
virtual ~CaffeNet() {}
inline void check_array_against_blob(
PyArrayObject* arr, Blob<float>* blob) {
CHECK(PyArray_FLAGS(arr) & NPY_ARRAY_C_CONTIGUOUS);
CHECK_EQ(PyArray_NDIM(arr), 4);
CHECK_EQ(PyArray_ITEMSIZE(arr), 4);
npy_intp* dims = PyArray_DIMS(arr);
CHECK_EQ(dims[0], blob->num());
CHECK_EQ(dims[1], blob->channels());
CHECK_EQ(dims[2], blob->height());
CHECK_EQ(dims[3], blob->width());
}
// The actual forward function. It takes in a python list of numpy arrays as
// input and a python list of numpy arrays as output. The input and output
// should all have correct shapes, are single-precisionabcdnt- and c contiguous.
void Forward(list bottom, list top) {
vector<Blob<float>*>& input_blobs = net_->input_blobs();
CHECK_EQ(len(bottom), input_blobs.size());
CHECK_EQ(len(top), net_->num_outputs());
// First, copy the input
for (int i = 0; i < input_blobs.size(); ++i) {
object elem = bottom[i];
PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(elem.ptr());
check_array_against_blob(arr, input_blobs[i]);
switch (Caffe::mode()) {
case Caffe::CPU:
memcpy(input_blobs[i]->mutable_cpu_data(), PyArray_DATA(arr),
sizeof(float) * input_blobs[i]->count());
break;
case Caffe::GPU:
cudaMemcpy(input_blobs[i]->mutable_gpu_data(), PyArray_DATA(arr),
sizeof(float) * input_blobs[i]->count(), cudaMemcpyHostToDevice);
break;
default:
LOG(FATAL) << "Unknown Caffe mode.";
} // switch (Caffe::mode())
}
//LOG(INFO) << "Start";
const vector<Blob<float>*>& output_blobs = net_->ForwardPrefilled();
//LOG(INFO) << "End";
for (int i = 0; i < output_blobs.size(); ++i) {
object elem = top[i];
PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(elem.ptr());
check_array_against_blob(arr, output_blobs[i]);
switch (Caffe::mode()) {
case Caffe::CPU:
memcpy(PyArray_DATA(arr), output_blobs[i]->cpu_data(),
sizeof(float) * output_blobs[i]->count());
break;
case Caffe::GPU:
cudaMemcpy(PyArray_DATA(arr), output_blobs[i]->gpu_data(),
sizeof(float) * output_blobs[i]->count(), cudaMemcpyDeviceToHost);
break;
default:
LOG(FATAL) << "Unknown Caffe mode.";
} // switch (Caffe::mode())
}
}
void Backward(list top_diff, list bottom_diff) {
vector<Blob<float>*>& output_blobs = net_->output_blobs();
vector<Blob<float>*>& input_blobs = net_->input_blobs();
CHECK_EQ(len(bottom_diff), input_blobs.size());
CHECK_EQ(len(top_diff), output_blobs.size());
// First, copy the output diff
for (int i = 0; i < output_blobs.size(); ++i) {
object elem = top_diff[i];
PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(elem.ptr());
check_array_against_blob(arr, output_blobs[i]);
switch (Caffe::mode()) {
case Caffe::CPU:
memcpy(output_blobs[i]->mutable_cpu_diff(), PyArray_DATA(arr),
sizeof(float) * output_blobs[i]->count());
break;
case Caffe::GPU:
cudaMemcpy(output_blobs[i]->mutable_gpu_diff(), PyArray_DATA(arr),
sizeof(float) * output_blobs[i]->count(), cudaMemcpyHostToDevice);
break;
default:
LOG(FATAL) << "Unknown Caffe mode.";
} // switch (Caffe::mode())
}
//LOG(INFO) << "Start";
net_->Backward();
//LOG(INFO) << "End";
for (int i = 0; i < input_blobs.size(); ++i) {
object elem = bottom_diff[i];
PyArrayObject* arr = reinterpret_cast<PyArrayObject*>(elem.ptr());
check_array_against_blob(arr, input_blobs[i]);
switch (Caffe::mode()) {
case Caffe::CPU:
memcpy(PyArray_DATA(arr), input_blobs[i]->cpu_diff(),
sizeof(float) * input_blobs[i]->count());
break;
case Caffe::GPU:
cudaMemcpy(PyArray_DATA(arr), input_blobs[i]->gpu_diff(),
sizeof(float) * input_blobs[i]->count(), cudaMemcpyDeviceToHost);
break;
default:
LOG(FATAL) << "Unknown Caffe mode.";
} // switch (Caffe::mode())
}
}
// The caffe::Caffe utility functions.
void set_mode_cpu() { Caffe::set_mode(Caffe::CPU); }
void set_mode_gpu() { Caffe::set_mode(Caffe::GPU); }
void set_phase_train() { Caffe::set_phase(Caffe::TRAIN); }
void set_phase_test() { Caffe::set_phase(Caffe::TEST); }
void set_device(int device_id) { Caffe::SetDevice(device_id); }
// The pointer to the internal caffe::Net instant.
shared_ptr<Net<float> > net_;
};
// The boost python module definition.
BOOST_PYTHON_MODULE(pycaffe)
{
boost::python::class_<CaffeNet>(
"CaffeNet", boost::python::init<string, string>())
.def("Forward", &CaffeNet::Forward)
.def("Backward", &CaffeNet::Backward)
.def("set_mode_cpu", &CaffeNet::set_mode_cpu)
.def("set_mode_gpu", &CaffeNet::set_mode_gpu)
.def("set_phase_train", &CaffeNet::set_phase_train)
.def("set_phase_test", &CaffeNet::set_phase_test)
.def("set_device", &CaffeNet::set_device)
;
}