mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
8c3a94eaf2
Summary: To illustrate the benefits of this commit, I'll use the time/iter I got from one of the JIT benchmarks on my machine. | Run | Time | |----------------------------------------------|-------------------------| | No profiler | 45ms | | With profiler | 56ms | | Use `clock_gettime` instead of `std::chrono` | 48ms | | Touch all pages on block allocation | 48ms (less jitter) | | Use `const char*` instead of `std::string` | 47ms (even less jitter) | Pull Request resolved: https://github.com/pytorch/pytorch/pull/11773 Differential Revision: D9886858 Pulled By: apaszke fbshipit-source-id: 58f926f09e95df0b11ec687763a72b06b66991d0
117 lines
3.9 KiB
C++
117 lines
3.9 KiB
C++
#include "torch/csrc/python_headers.h"
|
|
|
|
#include "torch/csrc/Exceptions.h"
|
|
#include "torch/csrc/utils/pybind.h"
|
|
#include "torch/csrc/autograd/grad_mode.h"
|
|
#include "torch/csrc/autograd/profiler.h"
|
|
#include "torch/csrc/autograd/python_function.h"
|
|
#include "torch/csrc/autograd/function.h"
|
|
|
|
PyObject * THPAutograd_initExtension(PyObject *_unused)
|
|
{
|
|
auto tensor_module = THPObjectPtr(PyImport_ImportModule("torch.tensor"));
|
|
if (!tensor_module) throw python_error();
|
|
|
|
// NOTE: "leaks" THPVariableClass
|
|
THPVariableClass = PyObject_GetAttrString(tensor_module, "Tensor");
|
|
if (!THPVariableClass) throw python_error();
|
|
|
|
auto autograd_module = THPObjectPtr(PyImport_ImportModule("torch.autograd"));
|
|
if (!autograd_module) throw python_error();
|
|
|
|
// NOTE: "leaks" Function
|
|
THPFunctionClass = PyObject_GetAttrString(autograd_module, "Function");
|
|
if (!THPFunctionClass) throw python_error();
|
|
|
|
auto m = py::handle(autograd_module).cast<py::module>();
|
|
|
|
py::class_<torch::autograd::profiler::Event>(m, "ProfilerEvent")
|
|
.def("kind", &torch::autograd::profiler::Event::kind)
|
|
.def(
|
|
"name",
|
|
[](const torch::autograd::profiler::Event& e) { return e.name(); })
|
|
.def("thread_id", &torch::autograd::profiler::Event::thread_id)
|
|
.def("device", &torch::autograd::profiler::Event::device)
|
|
.def("cpu_elapsed_us", &torch::autograd::profiler::Event::cpu_elapsed_us)
|
|
.def(
|
|
"cuda_elapsed_us", &torch::autograd::profiler::Event::cuda_elapsed_us)
|
|
.def("has_cuda", &torch::autograd::profiler::Event::has_cuda);
|
|
py::enum_<torch::autograd::profiler::ProfilerState>(m,"ProfilerState")
|
|
.value("Disabled", torch::autograd::profiler::ProfilerState::Disabled)
|
|
.value("CPU", torch::autograd::profiler::ProfilerState::CPU)
|
|
.value("CUDA", torch::autograd::profiler::ProfilerState::CUDA)
|
|
.value("NVTX", torch::autograd::profiler::ProfilerState::NVTX);
|
|
|
|
m.def("_enable_profiler", torch::autograd::profiler::enableProfiler);
|
|
m.def("_disable_profiler", torch::autograd::profiler::disableProfiler);
|
|
|
|
m.def("_push_range", [](std::string name) {
|
|
torch::autograd::profiler::pushRange(std::move(name));
|
|
});
|
|
m.def("_pop_range", []() { torch::autograd::profiler::popRange(); });
|
|
|
|
/// TODO: Replace this ASAP with a better solution for deep autograd graphs!
|
|
m.def("_unsafe_set_delete_function_max_recursion_depth", [](size_t value) {
|
|
torch::autograd::deleteFunctionMaxRecursionDepth = value;
|
|
});
|
|
|
|
Py_RETURN_TRUE;
|
|
}
|
|
|
|
namespace torch { namespace autograd {
|
|
|
|
static PyObject * set_grad_enabled(PyObject* _unused, PyObject *arg) {
|
|
HANDLE_TH_ERRORS
|
|
if (!PyBool_Check(arg)) {
|
|
throw TypeError("enabled must be a bool (got %s)", Py_TYPE(arg)->tp_name);
|
|
}
|
|
GradMode::set_enabled(arg == Py_True);
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject * is_grad_enabled(PyObject* _unused, PyObject *arg) {
|
|
HANDLE_TH_ERRORS
|
|
if (GradMode::is_enabled()) {
|
|
Py_RETURN_TRUE;
|
|
} else {
|
|
Py_RETURN_FALSE;
|
|
}
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject * set_anomaly_mode_enabled(PyObject* _unused, PyObject *arg) {
|
|
HANDLE_TH_ERRORS
|
|
if (!PyBool_Check(arg)) {
|
|
throw TypeError("enabled must be a bool (got %s)", Py_TYPE(arg)->tp_name);
|
|
}
|
|
AnomalyMode::set_enabled(arg == Py_True);
|
|
Py_RETURN_NONE;
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
static PyObject * is_anomaly_mode_enabled(PyObject* _unused, PyObject *arg) {
|
|
HANDLE_TH_ERRORS
|
|
if (AnomalyMode::is_enabled()) {
|
|
Py_RETURN_TRUE;
|
|
} else {
|
|
Py_RETURN_FALSE;
|
|
}
|
|
END_HANDLE_TH_ERRORS
|
|
}
|
|
|
|
// autograd methods on torch._C
|
|
static PyMethodDef methods[] = {
|
|
{"set_grad_enabled", (PyCFunction)set_grad_enabled, METH_O, nullptr},
|
|
{"is_grad_enabled", (PyCFunction)is_grad_enabled, METH_NOARGS, nullptr},
|
|
{"set_anomaly_enabled", (PyCFunction)set_anomaly_mode_enabled, METH_O, nullptr},
|
|
{"is_anomaly_enabled", (PyCFunction)is_anomaly_mode_enabled, METH_NOARGS, nullptr},
|
|
{nullptr, nullptr, 0, nullptr}
|
|
};
|
|
|
|
PyMethodDef* python_functions() {
|
|
return methods;
|
|
}
|
|
|
|
}} // namespace torch::autograd
|