comp-chem/cp2k
1
0
Fork 0
mirror of https://github.com/cp2k/cp2k.git synced 2026-07-28 22:25:32 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

Torch: Add autograd to Fortran API

Browse source
This commit is contained in:
Ole Schütt 2025-01-25 15:32:33 +01:00 • committed by Ole Schütt
parent 34d1e17cc8
commit b6963b94bc
2 changed files with 124 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

99
src/torch_api.F
View file

@ -65,7 +65,8 @@ MODULE torch_api
MODULE PROCEDURE torch_model_get_attr_strlist
END INTERFACE torch_model_get_attr
PUBLIC :: torch_tensor_type, torch_tensor_from_array, torch_tensor_data_ptr, torch_tensor_release
PUBLIC :: torch_tensor_type, torch_tensor_from_array, torch_tensor_release
PUBLIC :: torch_tensor_data_ptr, torch_tensor_backward, torch_tensor_grad
PUBLIC :: torch_dict_type, torch_dict_create, torch_dict_insert, torch_dict_get, torch_dict_release
PUBLIC :: torch_model_type, torch_model_load, torch_model_forward, torch_model_release
PUBLIC :: torch_model_get_attr, torch_model_read_metadata
@ -85,31 +86,38 @@ CONTAINS
!> The source must be an ALLOCATABLE to prevent passing a temporary array.
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_tensor_from_array_${typename}$_${ndims}$d(tensor, source)
SUBROUTINE torch_tensor_from_array_${typename}$_${ndims}$d(tensor, source, requires_grad)
TYPE(torch_tensor_type), INTENT(INOUT) :: tensor
#:set arraydims = ", ".join(":" for i in range(ndims))
${type_f}$, DIMENSION(${arraydims}$), ALLOCATABLE, INTENT(IN) :: source
LOGICAL, OPTIONAL, INTENT(IN) :: requires_grad
#if defined(__LIBTORCH)
INTEGER(kind=int_8), DIMENSION(${ndims}$) :: sizes_c
LOGICAL :: my_req_grad
INTERFACE
SUBROUTINE torch_c_tensor_from_array_${typename}$ (tensor, ndims, sizes, source) &
SUBROUTINE torch_c_tensor_from_array_${typename}$ (tensor, req_grad, ndims, sizes, source) &
BIND(C, name="torch_c_tensor_from_array_${typename}$")
IMPORT :: C_PTR, C_INT, C_INT64_T, C_FLOAT, C_DOUBLE
IMPORT :: C_PTR, C_INT, C_INT64_T, C_FLOAT, C_DOUBLE, C_BOOL
TYPE(C_PTR) :: tensor
LOGICAL(kind=C_BOOL), VALUE :: req_grad
INTEGER(kind=C_INT), VALUE :: ndims
INTEGER(kind=C_INT64_T), DIMENSION(*) :: sizes
${type_c}$, DIMENSION(*) :: source
END SUBROUTINE torch_c_tensor_from_array_${typename}$
END INTERFACE
my_req_grad = .FALSE.
IF (PRESENT(requires_grad)) my_req_grad = requires_grad
#:for axis in range(ndims)
sizes_c(${axis + 1}$) = SIZE(source, ${ndims - axis}$) ! C arrays are stored row-major.
#:endfor
CPASSERT(.NOT. C_ASSOCIATED(tensor%c_ptr))
CALL torch_c_tensor_from_array_${typename}$ (tensor=tensor%c_ptr, &
req_grad=LOGICAL(my_req_grad, C_BOOL), &
ndims=${ndims}$, &
sizes=sizes_c, &
source=source)
@ -174,6 +182,70 @@ CONTAINS
#:endfor
#:endfor
! **************************************************************************************************
!> \brief Runs autograd on a Torch tensor.
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_tensor_backward(tensor, outer_grad)
TYPE(torch_tensor_type), INTENT(IN) :: tensor
TYPE(torch_tensor_type), INTENT(IN) :: outer_grad
#if defined(__LIBTORCH)
CHARACTER(len=*), PARAMETER :: routineN = 'torch_tensor_backward'
INTEGER :: handle
INTERFACE
SUBROUTINE torch_c_tensor_backward(tensor, outer_grad) &
BIND(C, name="torch_c_tensor_backward")
IMPORT :: C_CHAR, C_PTR
TYPE(C_PTR), VALUE :: tensor
TYPE(C_PTR), VALUE :: outer_grad
END SUBROUTINE torch_c_tensor_backward
END INTERFACE
CALL timeset(routineN, handle)
CPASSERT(C_ASSOCIATED(tensor%c_ptr))
CPASSERT(C_ASSOCIATED(outer_grad%c_ptr))
CALL torch_c_tensor_backward(tensor=tensor%c_ptr, outer_grad=outer_grad%c_ptr)
CALL timestop(handle)
#else
CPABORT("CP2K compiled without the Torch library.")
MARK_USED(dict)
MARK_USED(key)
MARK_USED(tensor)
#endif
END SUBROUTINE torch_tensor_backward
! **************************************************************************************************
!> \brief Returns the gradient of a Torch tensor which was computed by autograd.
!> \author Ole Schuett
! **************************************************************************************************
SUBROUTINE torch_tensor_grad(tensor, grad)
TYPE(torch_tensor_type), INTENT(IN) :: tensor
TYPE(torch_tensor_type), INTENT(INOUT) :: grad
#if defined(__LIBTORCH)
INTERFACE
SUBROUTINE torch_c_tensor_grad(tensor, grad) &
BIND(C, name="torch_c_tensor_grad")
IMPORT :: C_PTR
TYPE(C_PTR), VALUE :: tensor
TYPE(C_PTR) :: grad
END SUBROUTINE torch_c_tensor_grad
END INTERFACE
CPASSERT(C_ASSOCIATED(tensor%c_ptr))
CPASSERT(.NOT. C_ASSOCIATED(grad%c_ptr))
CALL torch_c_tensor_grad(tensor=tensor%c_ptr, grad=grad%c_ptr)
CPASSERT(C_ASSOCIATED(grad%c_ptr))
#else
CPABORT("CP2K compiled without the Torch library.")
MARK_USED(dict)
MARK_USED(key)
MARK_USED(tensor)
#endif
END SUBROUTINE torch_tensor_grad
! **************************************************************************************************
!> \brief Releases a Torch tensor and all its ressources.
!> \author Ole Schuett
@ -321,6 +393,9 @@ CONTAINS
CHARACTER(len=*), INTENT(IN) :: filename
#if defined(__LIBTORCH)
CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_load'
INTEGER :: handle
INTERFACE
SUBROUTINE torch_c_model_load(model, filename) BIND(C, name="torch_c_model_load")
IMPORT :: C_PTR, C_CHAR
@ -329,9 +404,11 @@ CONTAINS
END SUBROUTINE torch_c_model_load
END INTERFACE
CALL timeset(routineN, handle)
CPASSERT(.NOT. C_ASSOCIATED(model%c_ptr))
CALL torch_c_model_load(model=model%c_ptr, filename=TRIM(filename)//C_NULL_CHAR)
CPASSERT(C_ASSOCIATED(model%c_ptr))
CALL timestop(handle)
#else
CPABORT("CP2K was compiled without Torch library.")
MARK_USED(model)
@ -349,6 +426,9 @@ CONTAINS
TYPE(torch_dict_type), INTENT(INOUT) :: outputs
#if defined(__LIBTORCH)
CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_forward'
INTEGER :: handle
INTERFACE
SUBROUTINE torch_c_model_forward(model, inputs, outputs) BIND(C, name="torch_c_model_forward")
IMPORT :: C_PTR
@ -358,12 +438,12 @@ CONTAINS
END SUBROUTINE torch_c_model_forward
END INTERFACE
CALL timeset(routineN, handle)
CPASSERT(C_ASSOCIATED(model%c_ptr))
CPASSERT(C_ASSOCIATED(inputs%c_ptr))
CPASSERT(C_ASSOCIATED(outputs%c_ptr))
CALL torch_c_model_forward(model=model%c_ptr, &
inputs=inputs%c_ptr, &
outputs=outputs%c_ptr)
CALL torch_c_model_forward(model=model%c_ptr, inputs=inputs%c_ptr, outputs=outputs%c_ptr)
CALL timestop(handle)
#else
CPABORT("CP2K was compiled without Torch library.")
MARK_USED(model)
@ -503,6 +583,9 @@ CONTAINS
TYPE(torch_model_type), INTENT(INOUT) :: model
#if defined(__LIBTORCH)
CHARACTER(len=*), PARAMETER :: routineN = 'torch_model_freeze'
INTEGER :: handle
INTERFACE
SUBROUTINE torch_c_model_freeze(model) BIND(C, name="torch_c_model_freeze")
IMPORT :: C_PTR
@ -510,8 +593,10 @@ CONTAINS
END SUBROUTINE torch_c_model_freeze
END INTERFACE
CALL timeset(routineN, handle)
CPASSERT(C_ASSOCIATED(model%c_ptr))
CALL torch_c_model_freeze(model=model%c_ptr)
CALL timestop(handle)
#else
CPABORT("CP2K was compiled without Torch library.")
MARK_USED(model)

42
src/torch_c_api.cpp
View file

@ -27,12 +27,12 @@ static torch::Device get_device() {
* \author Ole Schuett
******************************************************************************/
static torch_c_tensor_t *tensor_from_array(const torch::Dtype dtype,
const int ndims,
const bool req_grad, const int ndims,
const int64_t sizes[],
void *source) {
const auto options = torch::TensorOptions().dtype(dtype);
const auto opts = torch::TensorOptions().dtype(dtype).requires_grad(req_grad);
const auto sizes_ref = c10::IntArrayRef(sizes, ndims);
return new torch_c_tensor_t(torch::from_blob(source, sizes_ref, options));
return new torch_c_tensor_t(torch::from_blob(source, sizes_ref, opts));
}
/*******************************************************************************
@ -61,9 +61,10 @@ extern "C" {
* The passed array has to outlive the tensor!
* \author Ole Schuett
******************************************************************************/
void torch_c_tensor_from_array_float(torch_c_tensor_t **tensor, const int ndims,
void torch_c_tensor_from_array_float(torch_c_tensor_t **tensor,
const bool req_grad, const int ndims,
const int64_t sizes[], float source[]) {
*tensor = tensor_from_array(torch::kFloat32, ndims, sizes, source);
*tensor = tensor_from_array(torch::kFloat32, req_grad, ndims, sizes, source);
}
/*******************************************************************************
@ -71,9 +72,10 @@ void torch_c_tensor_from_array_float(torch_c_tensor_t **tensor, const int ndims,
* The passed array has to outlive the tensor!
* \author Ole Schuett
******************************************************************************/
void torch_c_tensor_from_array_int64(torch_c_tensor_t **tensor, const int ndims,
void torch_c_tensor_from_array_int64(torch_c_tensor_t **tensor,
const bool req_grad, const int ndims,
const int64_t sizes[], int64_t source[]) {
*tensor = tensor_from_array(torch::kInt64, ndims, sizes, source);
*tensor = tensor_from_array(torch::kInt64, req_grad, ndims, sizes, source);
}
/*******************************************************************************
@ -82,9 +84,9 @@ void torch_c_tensor_from_array_int64(torch_c_tensor_t **tensor, const int ndims,
* \author Ole Schuett
******************************************************************************/
void torch_c_tensor_from_array_double(torch_c_tensor_t **tensor,
const int ndims, const int64_t sizes[],
double source[]) {
*tensor = tensor_from_array(torch::kFloat64, ndims, sizes, source);
const bool req_grad, const int ndims,
const int64_t sizes[], double source[]) {
*tensor = tensor_from_array(torch::kFloat64, req_grad, ndims, sizes, source);
}
/*******************************************************************************
@ -120,6 +122,26 @@ void torch_c_tensor_data_ptr_double(const torch_c_tensor_t *tensor,
*data_ptr = (double *)get_data_ptr(tensor, torch::kFloat64, ndims, sizes);
}
/*******************************************************************************
* \brief Runs autograd on a Torch tensor.
* \author Ole Schuett
******************************************************************************/
void torch_c_tensor_backward(const torch_c_tensor_t *tensor,
const torch_c_tensor_t *outer_grad) {
tensor->backward(*outer_grad);
}
/*******************************************************************************
* \brief Returns the gradient of a Torch tensor which was computed by autograd.
* \author Ole Schuett
******************************************************************************/
void torch_c_tensor_grad(const torch_c_tensor_t *tensor,
torch_c_tensor_t **grad) {
const torch::Tensor maybe_grad = tensor->grad();
assert(maybe_grad.defined());
*grad = new torch_c_tensor_t(maybe_grad.cpu().contiguous());
}
/*******************************************************************************
* \brief Releases a Torch tensor and all its ressources.
* \author Ole Schuett