LCOV - code coverage report
Current view: top level - src - pao_ml_neuralnet.F (source / functions) Hit Total Coverage
Test: CP2K Regtests (git:0de0cc2) Lines: 107 108 99.1 %
Date: 2024-03-28 07:31:50 Functions: 5 5 100.0 %

          Line data    Source code
       1             : !--------------------------------------------------------------------------------------------------!
       2             : !   CP2K: A general program to perform molecular dynamics simulations                              !
       3             : !   Copyright 2000-2024 CP2K developers group <https://cp2k.org>                                   !
       4             : !                                                                                                  !
       5             : !   SPDX-License-Identifier: GPL-2.0-or-later                                                      !
       6             : !--------------------------------------------------------------------------------------------------!
       7             : 
       8             : ! **************************************************************************************************
       9             : !> \brief Neural Network implementation
      10             : !> \author Ole Schuett
      11             : ! **************************************************************************************************
      12             : MODULE pao_ml_neuralnet
      13             :    USE kinds,                           ONLY: dp
      14             :    USE pao_types,                       ONLY: pao_env_type,&
      15             :                                               training_matrix_type
      16             :    USE parallel_rng_types,              ONLY: rng_stream_type
      17             : #include "./base/base_uses.f90"
      18             : 
      19             :    IMPLICIT NONE
      20             : 
      21             :    PRIVATE
      22             : 
      23             :    CHARACTER(len=*), PARAMETER, PRIVATE :: moduleN = 'pao_ml_neuralnet'
      24             : 
      25             :    PUBLIC ::pao_ml_nn_train, pao_ml_nn_predict, pao_ml_nn_gradient
      26             : 
      27             :    ! TODO turn these into input parameters
      28             :    REAL(dp), PARAMETER   :: step_size = 0.001_dp
      29             :    INTEGER, PARAMETER    :: nlayers = 3
      30             :    REAL(dp), PARAMETER   :: convergence_eps = 1e-7_dp
      31             :    INTEGER, PARAMETER    :: max_training_cycles = 50000
      32             : 
      33             : CONTAINS
      34             : 
      35             : ! **************************************************************************************************
      36             : !> \brief Uses neural network to make a prediction
      37             : !> \param pao ...
      38             : !> \param ikind ...
      39             : !> \param descriptor ...
      40             : !> \param output ...
      41             : !> \param variance ...
      42             : ! **************************************************************************************************
      43          28 :    SUBROUTINE pao_ml_nn_predict(pao, ikind, descriptor, output, variance)
      44             :       TYPE(pao_env_type), POINTER                        :: pao
      45             :       INTEGER, INTENT(IN)                                :: ikind
      46             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: descriptor
      47             :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: output
      48             :       REAL(dp), INTENT(OUT)                              :: variance
      49             : 
      50             :       TYPE(training_matrix_type), POINTER                :: training_matrix
      51             : 
      52          28 :       training_matrix => pao%ml_training_matrices(ikind)
      53             : 
      54          28 :       CALL nn_eval(training_matrix%NN, input=descriptor, prediction=output)
      55             : 
      56          28 :       variance = 0.0_dp ! Neural Networks don't provide a variance
      57          28 :    END SUBROUTINE pao_ml_nn_predict
      58             : 
      59             : ! **************************************************************************************************
      60             : !> \brief Calculate gradient of neural network
      61             : !> \param pao ...
      62             : !> \param ikind ...
      63             : !> \param descriptor ...
      64             : !> \param outer_deriv ...
      65             : !> \param gradient ...
      66             : ! **************************************************************************************************
      67           4 :    SUBROUTINE pao_ml_nn_gradient(pao, ikind, descriptor, outer_deriv, gradient)
      68             :       TYPE(pao_env_type), POINTER                        :: pao
      69             :       INTEGER, INTENT(IN)                                :: ikind
      70             :       REAL(dp), DIMENSION(:), INTENT(IN), TARGET         :: descriptor
      71             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: outer_deriv
      72             :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: gradient
      73             : 
      74             :       INTEGER                                            :: i, ilayer, j, nlayers, width, width_in, &
      75             :                                                             width_out
      76           4 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: backward, forward
      77           4 :       REAL(dp), DIMENSION(:, :, :), POINTER              :: A
      78             : 
      79           4 :       A => pao%ml_training_matrices(ikind)%NN
      80             : 
      81           4 :       nlayers = SIZE(A, 1)
      82           0 :       width = SIZE(A, 2); CPASSERT(SIZE(A, 2) == SIZE(A, 3))
      83           4 :       width_in = SIZE(descriptor)
      84           4 :       width_out = SIZE(outer_deriv)
      85             : 
      86          28 :       ALLOCATE (forward(0:nlayers, width), backward(0:nlayers, width))
      87             : 
      88         144 :       forward = 0.0_dp
      89          16 :       forward(0, 1:width_in) = descriptor
      90             : 
      91          16 :       DO ilayer = 1, nlayers
      92         100 :       DO i = 1, width
      93         684 :       DO j = 1, width
      94         672 :          forward(ilayer, i) = forward(ilayer, i) + A(ilayer, i, j)*TANH(forward(ilayer - 1, j))
      95             :       END DO
      96             :       END DO
      97             :       END DO
      98             : 
      99             :       ! Turning Point ------------------------------------------------------------------------------
     100         144 :       backward = 0.0_dp
     101          32 :       backward(nlayers, 1:width_out) = outer_deriv(:)
     102             : 
     103          16 :       DO ilayer = nlayers, 1, -1
     104         100 :       DO i = 1, width
     105         684 :       DO j = 1, width
     106         672 :   backward(ilayer - 1, j) = backward(ilayer - 1, j) + backward(ilayer, i)*A(ilayer, i, j)*(1.0_dp - TANH(forward(ilayer - 1, j))**2)
     107             :       END DO
     108             :       END DO
     109             :       END DO
     110             : 
     111          16 :       gradient(:) = backward(0, 1:width_in)
     112             : 
     113           4 :       DEALLOCATE (forward, backward)
     114           4 :    END SUBROUTINE pao_ml_nn_gradient
     115             : 
     116             : ! **************************************************************************************************
     117             : !> \brief Trains the neural network on given training points
     118             : !> \param pao ...
     119             : ! **************************************************************************************************
     120           4 :    SUBROUTINE pao_ml_nn_train(pao)
     121             :       TYPE(pao_env_type), POINTER                        :: pao
     122             : 
     123             :       INTEGER                                            :: i, icycle, ikind, ilayer, ipoint, j, &
     124             :                                                             npoints, width, width_in, width_out
     125             :       REAL(dp)                                           :: bak, eps, error, error1, error2, num_grad
     126           4 :       REAL(dp), ALLOCATABLE, DIMENSION(:)                :: prediction
     127           4 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :, :)          :: gradient
     128             :       TYPE(rng_stream_type)                              :: rng_stream
     129             :       TYPE(training_matrix_type), POINTER                :: training_matrix
     130             : 
     131             :       ! TODO this could be parallelized over ranks
     132           8 :       DO ikind = 1, SIZE(pao%ml_training_matrices)
     133           4 :          training_matrix => pao%ml_training_matrices(ikind)
     134             : 
     135           4 :          npoints = SIZE(training_matrix%inputs, 2) ! number of points
     136           4 :          CPASSERT(SIZE(training_matrix%outputs, 2) == npoints)
     137           4 :          IF (npoints == 0) CYCLE
     138             : 
     139             :          !TODO proper output
     140           6 :          IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| Training neural network for kind: ", &
     141           4 :             TRIM(training_matrix%kindname), " from ", npoints, "training points."
     142             : 
     143             :          ! determine network width and allocate it
     144           4 :          width_in = SIZE(training_matrix%inputs, 1)
     145           4 :          width_out = SIZE(training_matrix%outputs, 1)
     146           4 :          width = MAX(width_in, width_out)
     147          16 :          ALLOCATE (training_matrix%NN(nlayers, width, width))
     148             : 
     149             :          ! initialize network with random numbers from -1.0 ... +1.0
     150           4 :          rng_stream = rng_stream_type(name="pao_nn")
     151          16 :          DO ilayer = 1, nlayers
     152         100 :          DO i = 1, width
     153         684 :          DO j = 1, width
     154         672 :             training_matrix%NN(ilayer, i, j) = -1.0_dp + 2.0_dp*rng_stream%next()
     155             :          END DO
     156             :          END DO
     157             :          END DO
     158             : 
     159             :          ! train the network using backpropagation
     160          16 :          ALLOCATE (gradient(nlayers, width, width))
     161      183452 :          DO icycle = 1, max_training_cycles
     162      183450 :             error = 0.0_dp
     163    37423800 :             gradient = 0.0_dp
     164      917250 :             DO ipoint = 1, npoints
     165             :                CALL nn_backpropagate(training_matrix%NN, &
     166             :                                      input=training_matrix%inputs(:, ipoint), &
     167             :                                      goal=training_matrix%outputs(:, ipoint), &
     168             :                                      gradient=gradient, &
     169      917250 :                                      error=error)
     170             :             END DO
     171    37423800 :             training_matrix%NN(:, :, :) = training_matrix%NN - step_size*gradient
     172             : 
     173      183450 :             IF (pao%iw > 0 .AND. MOD(icycle, 100) == 0) WRITE (pao%iw, *) &
     174         917 :                "PAO|ML| ", TRIM(training_matrix%kindname), &
     175      187985 :                " training-cycle:", icycle, "SQRT(error):", SQRT(error), "grad:", SUM(gradient**2)
     176             : 
     177    37423802 :             IF (SUM(gradient**2) < convergence_eps) EXIT
     178             :          END DO
     179             : 
     180             :          ! numeric gradient for debugging ----------------------------------------------------------
     181             :          IF (.FALSE.) THEN
     182             :             eps = 1e-4_dp
     183             :             ilayer = 1
     184             :             ipoint = 1
     185             :             error = 0.0_dp
     186             :             gradient = 0.0_dp
     187             :             CALL nn_backpropagate(training_matrix%NN, &
     188             :                                   input=training_matrix%inputs(:, ipoint), &
     189             :                                   goal=training_matrix%outputs(:, ipoint), &
     190             :                                   gradient=gradient, &
     191             :                                   error=error)
     192             : 
     193             :             ALLOCATE (prediction(width_out))
     194             :             DO i = 1, width
     195             :             DO j = 1, width
     196             :                bak = training_matrix%NN(ilayer, i, j)
     197             : 
     198             :                training_matrix%NN(ilayer, i, j) = bak + eps
     199             :                CALL nn_eval(training_matrix%NN, &
     200             :                             input=training_matrix%inputs(:, ipoint), &
     201             :                             prediction=prediction)
     202             :                error1 = SUM((training_matrix%outputs(:, ipoint) - prediction)**2)
     203             : 
     204             :                training_matrix%NN(ilayer, i, j) = bak - eps
     205             :                CALL nn_eval(training_matrix%NN, &
     206             :                             input=training_matrix%inputs(:, ipoint), &
     207             :                             prediction=prediction)
     208             :                error2 = SUM((training_matrix%outputs(:, ipoint) - prediction)**2)
     209             : 
     210             :                training_matrix%NN(ilayer, i, j) = bak
     211             :                num_grad = (error1 - error2)/(2.0_dp*eps)
     212             :                IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| Numeric gradient:", i, j, gradient(ilayer, i, j), num_grad
     213             : 
     214             :             END DO
     215             :             END DO
     216             :             DEALLOCATE (prediction)
     217             :          END IF
     218             :          !------------------------------------------------------------------------------------------
     219             : 
     220           4 :          DEALLOCATE (gradient)
     221             : 
     222             :          ! test training points individually
     223          12 :          ALLOCATE (prediction(width_out))
     224          20 :          DO ipoint = 1, npoints
     225             :             CALL nn_eval(training_matrix%NN, &
     226             :                          input=training_matrix%inputs(:, ipoint), &
     227          16 :                          prediction=prediction)
     228         144 :             error = MAXVAL(ABS(training_matrix%outputs(:, ipoint) - prediction))
     229          24 :             IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| ", TRIM(training_matrix%kindname), &
     230          20 :                " verify training-point:", ipoint, "SQRT(error):", SQRT(error)
     231             :          END DO
     232           8 :          DEALLOCATE (prediction)
     233             : 
     234             :       END DO
     235             : 
     236         104 :    END SUBROUTINE pao_ml_nn_train
     237             : 
     238             : ! **************************************************************************************************
     239             : !> \brief Evaluates the neural network for a given input
     240             : !> \param A ...
     241             : !> \param input ...
     242             : !> \param prediction ...
     243             : ! **************************************************************************************************
     244          44 :    SUBROUTINE nn_eval(A, input, prediction)
     245             :       REAL(dp), DIMENSION(:, :, :), INTENT(IN)           :: A
     246             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: input
     247             :       REAL(dp), DIMENSION(:), INTENT(OUT)                :: prediction
     248             : 
     249             :       INTEGER                                            :: i, ilayer, j, nlayers, width, width_in, &
     250             :                                                             width_out
     251          44 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: forward
     252             : 
     253          44 :       nlayers = SIZE(A, 1)
     254          44 :       width = SIZE(A, 2); CPASSERT(SIZE(A, 2) == SIZE(A, 3))
     255          44 :       width_in = SIZE(input)
     256          44 :       width_out = SIZE(prediction)
     257             : 
     258         176 :       ALLOCATE (forward(0:nlayers, width))
     259             : 
     260        1584 :       forward = 0.0_dp
     261         128 :       forward(0, 1:width_in) = input(:)
     262             : 
     263         176 :       DO ilayer = 1, nlayers
     264        1100 :       DO i = 1, width
     265        7524 :       DO j = 1, width
     266        7392 :          forward(ilayer, i) = forward(ilayer, i) + A(ilayer, i, j)*TANH(forward(ilayer - 1, j))
     267             :       END DO
     268             :       END DO
     269             :       END DO
     270             : 
     271         352 :       prediction(:) = forward(nlayers, 1:width_out)
     272             : 
     273          44 :    END SUBROUTINE nn_eval
     274             : 
     275             : ! **************************************************************************************************
     276             : !> \brief Uses backpropagation to calculate the gradient for a given training point
     277             : !> \param A ...
     278             : !> \param input ...
     279             : !> \param goal ...
     280             : !> \param error ...
     281             : !> \param gradient ...
     282             : ! **************************************************************************************************
     283      733800 :    SUBROUTINE nn_backpropagate(A, input, goal, error, gradient)
     284             :       REAL(dp), DIMENSION(:, :, :), INTENT(IN)           :: A
     285             :       REAL(dp), DIMENSION(:), INTENT(IN)                 :: input, goal
     286             :       REAL(dp), INTENT(INOUT)                            :: error
     287             :       REAL(dp), DIMENSION(:, :, :), INTENT(INOUT)        :: gradient
     288             : 
     289             :       INTEGER                                            :: i, ilayer, j, nlayers, width, width_in, &
     290             :                                                             width_out
     291      733800 :       REAL(dp), ALLOCATABLE, DIMENSION(:)                :: prediction
     292      733800 :       REAL(dp), ALLOCATABLE, DIMENSION(:, :)             :: backward, forward
     293             : 
     294      733800 :       nlayers = SIZE(A, 1)
     295      733800 :       width = SIZE(A, 2); CPASSERT(SIZE(A, 2) == SIZE(A, 3))
     296      733800 :       width_in = SIZE(input)
     297      733800 :       width_out = SIZE(goal)
     298             : 
     299     6604200 :       ALLOCATE (forward(0:nlayers, width), prediction(width_out), backward(0:nlayers, width))
     300             : 
     301    26416800 :       forward = 0.0_dp
     302     2802800 :       forward(0, 1:width_in) = input
     303             : 
     304     2935200 :       DO ilayer = 1, nlayers
     305    18345000 :       DO i = 1, width
     306   125479800 :       DO j = 1, width
     307   123278400 :          forward(ilayer, i) = forward(ilayer, i) + A(ilayer, i, j)*TANH(forward(ilayer - 1, j))
     308             :       END DO
     309             :       END DO
     310             :       END DO
     311             : 
     312     5870400 :       prediction(:) = forward(nlayers, 1:width_out)
     313             : 
     314     5870400 :       error = error + SUM((prediction - goal)**2)
     315             : 
     316             :       ! Turning Point ------------------------------------------------------------------------------
     317    26416800 :       backward = 0.0_dp
     318     5870400 :       backward(nlayers, 1:width_out) = prediction - goal
     319             : 
     320     2935200 :       DO ilayer = nlayers, 1, -1
     321    18345000 :       DO i = 1, width
     322   125479800 :       DO j = 1, width
     323   107868600 :          gradient(ilayer, i, j) = gradient(ilayer, i, j) + 2.0_dp*backward(ilayer, i)*TANH(forward(ilayer - 1, j))
     324   123278400 :   backward(ilayer - 1, j) = backward(ilayer - 1, j) + backward(ilayer, i)*A(ilayer, i, j)*(1.0_dp - TANH(forward(ilayer - 1, j))**2)
     325             :       END DO
     326             :       END DO
     327             :       END DO
     328             : 
     329      733800 :       DEALLOCATE (forward, backward, prediction)
     330      733800 :    END SUBROUTINE nn_backpropagate
     331             : 
     332             : END MODULE pao_ml_neuralnet

Generated by: LCOV version 1.15