python/cpp_bindings/tensor_py.cpp - with_cpu - vuron/adept

python/cpp_bindings/tensor_py.cpp
135 строк6 KB
kolkir
Fix memroy acces in batchnorm, update resnet training script
08 мар 2025, 13:46
08 мар 2025, 13:461476fed
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <adept/exception.hpp>
#include <adept/index_utils.hpp>
#include <adept/tensor.hpp>
#include <adept/tensor_print.hpp>
#include <adept/types_dispatch.hpp>

#include <sstream>

namespace py = pybind11;

using namespace adept;

void bind_tensor(py::module& m) {
  py::class_<Tensor>(m, "Tensor")
      .def(py::init([](const std::vector<float32_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Float32)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        return Tensor::from_blob(data.data(), props);
      }))
      .def(py::init([](const std::vector<float64_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Float64)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");

        return Tensor::from_blob(data.data(), props);
      }))
      .def(py::init([](const std::vector<int32_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Int32)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        return Tensor::from_blob(data.data(), props);
      }))
      .def(py::init([](const std::vector<int8_t>& data, const TensorProperties& props) {
        if (props.dtype != dtype_t::Int8)
          THROW_ERROR("Inconsistent type ", props.dtype, " for Tensor creation!");
        return Tensor::from_blob(data.data(), props);
      }))
      .def_property_readonly_static("empty",
                                    [](py::object) {
                                      return py::cpp_function([](const TensorProperties& props) {
                                        return Tensor::empty(props);
                                      });
                                    })
      .def_property_readonly_static("zero",
                                    [](py::object) {
                                      return py::cpp_function([](const TensorProperties& props) {
                                        return Tensor::zero(props);
                                      });
                                    })
      .def_property_readonly_static("uniform",
                                    [](py::object) {
                                      return py::cpp_function([](float32_t start, float32_t end,
                                                                 const TensorProperties& props) {
                                        return Tensor::uniform(start, end, props);
                                      });
                                    })
      .def_property_readonly_static("normal",
                                    [](py::object) {
                                      return py::cpp_function([](float32_t mean, float32_t std,
                                                                 const TensorProperties& props) {
                                        return Tensor::normal(mean, std, props);
                                      });
                                    })
      .def_property_readonly_static(
          "stack",
          [](py::object) {
            return py::cpp_function(
                [](const std::vector<Tensor>& tensors) { return Tensor::stack(tensors); });
          })
      .def(py::init([](py::buffer b) {
        py::buffer_info info = b.request();
        dtype_t dtype = dtype_t::Float32;
        if (info.item_type_is_equivalent_to<float32_t>()) {
          CHECK(sizeof(float32_t) == info.itemsize, "Invlaid python buffer type size");
          dtype = dtype_t::Float32;
        } else if (info.item_type_is_equivalent_to<float64_t>()) {
          CHECK(sizeof(float64_t) == info.itemsize, "Invlaid python buffer type size");
          dtype = dtype_t::Float64;
        } else if (info.item_type_is_equivalent_to<int32_t>()) {
          CHECK(sizeof(int32_t) == info.itemsize, "Invlaid python buffer type size");
          dtype = dtype_t::Int32;
        } else {
          THROW_ERROR("Tensor can't be constructed with the ", info.format, " type");
        }
        TensorProperties props(Shape(info.shape), device_t::CPU, dtype);
        return Tensor::from_blob(info.ptr, props);
      }))
      .def("numpy",
           [](Tensor& t) -> py::buffer_info {
             if (t.device() != device_t::CPU) {
               THROW_ERROR("Can't convert GPU tensor to numpy");
             }
             py::buffer_info info;
             DISPATCH_TYPE(t.dtype(), [&]() {
               info = py::buffer_info(t.mutable_data_ptr<scalar_t>(),     /* Pointer to buffer */
                                      t.shape().dims(),                   /* shape_in */
                                      make_continuous_strides(t.shape()), /* stride_in */
                                      true /* Readonly */);
             });
             return info;
           })
      .def("clone", &Tensor::clone)
      .def("properties", &Tensor::properties)
      .def("sum", &Tensor::sum)
      .def("mean", &Tensor::mean)
      .def("shape", [](const Tensor& self) { return self.properties().shape; })
      .def("__repr__",
           [](const Tensor& t) {
             std::stringstream buf;
             buf << t;
             return buf.str();
           })
      .def("__add__", py::overload_cast<const Tensor&>(&Tensor::add_))
      .def("__add__", py::overload_cast<float32_t>(&Tensor::add_))
      .def("__sub__", py::overload_cast<const Tensor&>(&Tensor::sub_))
      .def("__sub__", py::overload_cast<float32_t>(&Tensor::sub_))
      .def("__mul__", py::overload_cast<const Tensor&>(&Tensor::mul_))
      .def("__mul__", py::overload_cast<float32_t>(&Tensor::mul_))
      .def("__truediv__", py::overload_cast<const Tensor&>(&Tensor::div_))
      .def("__truediv__", py::overload_cast<float32_t>(&Tensor::div_))
      .def("__gt__", &Tensor::gt_)
      .def("__ge__", &Tensor::ge_)
      .def("__lt__", &Tensor::lt_)
      .def("__le__", &Tensor::le_)
      .def("__neg__", &Tensor::neg_)
      .def("exp_", &Tensor::exp_)
      .def("sqrt_", &Tensor::sqrt_)
      .def("sqrt", &Tensor::sqrt)
      .def("transpose2d", &Tensor::transpose2d)
      .def("matmul", &Tensor::matmul)
      .def("float_at", [](const Tensor& self, const coords_t& coords) -> float32_t {
        return self.at<float32_t>(coords);
      });
}