(git:ccc2433)
pao_ml_neuralnet.F
Go to the documentation of this file.
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  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  training_matrix => pao%ml_training_matrices(ikind)
53 
54  CALL nn_eval(training_matrix%NN, input=descriptor, prediction=output)
55 
56  variance = 0.0_dp ! Neural Networks don't provide a variance
57  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  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  REAL(dp), ALLOCATABLE, DIMENSION(:, :) :: backward, forward
77  REAL(dp), DIMENSION(:, :, :), POINTER :: a
78 
79  a => pao%ml_training_matrices(ikind)%NN
80 
81  nlayers = SIZE(a, 1)
82  width = SIZE(a, 2); cpassert(SIZE(a, 2) == SIZE(a, 3))
83  width_in = SIZE(descriptor)
84  width_out = SIZE(outer_deriv)
85 
86  ALLOCATE (forward(0:nlayers, width), backward(0:nlayers, width))
87 
88  forward = 0.0_dp
89  forward(0, 1:width_in) = descriptor
90 
91  DO ilayer = 1, nlayers
92  DO i = 1, width
93  DO j = 1, width
94  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  backward = 0.0_dp
101  backward(nlayers, 1:width_out) = outer_deriv(:)
102 
103  DO ilayer = nlayers, 1, -1
104  DO i = 1, width
105  DO j = 1, width
106  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  gradient(:) = backward(0, 1:width_in)
112 
113  DEALLOCATE (forward, backward)
114  END SUBROUTINE pao_ml_nn_gradient
115 
116 ! **************************************************************************************************
117 !> \brief Trains the neural network on given training points
118 !> \param pao ...
119 ! **************************************************************************************************
120  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  REAL(dp), ALLOCATABLE, DIMENSION(:) :: prediction
127  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  DO ikind = 1, SIZE(pao%ml_training_matrices)
133  training_matrix => pao%ml_training_matrices(ikind)
134 
135  npoints = SIZE(training_matrix%inputs, 2) ! number of points
136  cpassert(SIZE(training_matrix%outputs, 2) == npoints)
137  IF (npoints == 0) cycle
138 
139  !TODO proper output
140  IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| Training neural network for kind: ", &
141  trim(training_matrix%kindname), " from ", npoints, "training points."
142 
143  ! determine network width and allocate it
144  width_in = SIZE(training_matrix%inputs, 1)
145  width_out = SIZE(training_matrix%outputs, 1)
146  width = max(width_in, width_out)
147  ALLOCATE (training_matrix%NN(nlayers, width, width))
148 
149  ! initialize network with random numbers from -1.0 ... +1.0
150  rng_stream = rng_stream_type(name="pao_nn")
151  DO ilayer = 1, nlayers
152  DO i = 1, width
153  DO j = 1, width
154  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  ALLOCATE (gradient(nlayers, width, width))
161  DO icycle = 1, max_training_cycles
162  error = 0.0_dp
163  gradient = 0.0_dp
164  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  error=error)
170  END DO
171  training_matrix%NN(:, :, :) = training_matrix%NN - step_size*gradient
172 
173  IF (pao%iw > 0 .AND. mod(icycle, 100) == 0) WRITE (pao%iw, *) &
174  "PAO|ML| ", trim(training_matrix%kindname), &
175  " training-cycle:", icycle, "SQRT(error):", sqrt(error), "grad:", sum(gradient**2)
176 
177  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  DEALLOCATE (gradient)
221 
222  ! test training points individually
223  ALLOCATE (prediction(width_out))
224  DO ipoint = 1, npoints
225  CALL nn_eval(training_matrix%NN, &
226  input=training_matrix%inputs(:, ipoint), &
227  prediction=prediction)
228  error = maxval(abs(training_matrix%outputs(:, ipoint) - prediction))
229  IF (pao%iw > 0) WRITE (pao%iw, *) "PAO|ML| ", trim(training_matrix%kindname), &
230  " verify training-point:", ipoint, "SQRT(error):", sqrt(error)
231  END DO
232  DEALLOCATE (prediction)
233 
234  END DO
235 
236  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  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  REAL(dp), ALLOCATABLE, DIMENSION(:, :) :: forward
252 
253  nlayers = SIZE(a, 1)
254  width = SIZE(a, 2); cpassert(SIZE(a, 2) == SIZE(a, 3))
255  width_in = SIZE(input)
256  width_out = SIZE(prediction)
257 
258  ALLOCATE (forward(0:nlayers, width))
259 
260  forward = 0.0_dp
261  forward(0, 1:width_in) = input(:)
262 
263  DO ilayer = 1, nlayers
264  DO i = 1, width
265  DO j = 1, width
266  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  prediction(:) = forward(nlayers, 1:width_out)
272 
273  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  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  REAL(dp), ALLOCATABLE, DIMENSION(:) :: prediction
292  REAL(dp), ALLOCATABLE, DIMENSION(:, :) :: backward, forward
293 
294  nlayers = SIZE(a, 1)
295  width = SIZE(a, 2); cpassert(SIZE(a, 2) == SIZE(a, 3))
296  width_in = SIZE(input)
297  width_out = SIZE(goal)
298 
299  ALLOCATE (forward(0:nlayers, width), prediction(width_out), backward(0:nlayers, width))
300 
301  forward = 0.0_dp
302  forward(0, 1:width_in) = input
303 
304  DO ilayer = 1, nlayers
305  DO i = 1, width
306  DO j = 1, width
307  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  prediction(:) = forward(nlayers, 1:width_out)
313 
314  error = error + sum((prediction - goal)**2)
315 
316  ! Turning Point ------------------------------------------------------------------------------
317  backward = 0.0_dp
318  backward(nlayers, 1:width_out) = prediction - goal
319 
320  DO ilayer = nlayers, 1, -1
321  DO i = 1, width
322  DO j = 1, width
323  gradient(ilayer, i, j) = gradient(ilayer, i, j) + 2.0_dp*backward(ilayer, i)*tanh(forward(ilayer - 1, j))
324  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  DEALLOCATE (forward, backward, prediction)
330  END SUBROUTINE nn_backpropagate
331 
332 END MODULE pao_ml_neuralnet
kinds
Defines the basic variable types.
Definition: kinds.F:23
kinds::dp
integer, parameter, public dp
Definition: kinds.F:34
pao_ml_neuralnet
Neural Network implementation.
Definition: pao_ml_neuralnet.F:12
pao_ml_neuralnet::pao_ml_nn_gradient
subroutine, public pao_ml_nn_gradient(pao, ikind, descriptor, outer_deriv, gradient)
Calculate gradient of neural network.
Definition: pao_ml_neuralnet.F:68
pao_ml_neuralnet::pao_ml_nn_train
subroutine, public pao_ml_nn_train(pao)
Trains the neural network on given training points.
Definition: pao_ml_neuralnet.F:121
pao_ml_neuralnet::pao_ml_nn_predict
subroutine, public pao_ml_nn_predict(pao, ikind, descriptor, output, variance)
Uses neural network to make a prediction.
Definition: pao_ml_neuralnet.F:44
pao_types
Types used by the PAO machinery.
Definition: pao_types.F:12
parallel_rng_types
Parallel (pseudo)random number generator (RNG) for multiple streams and substreams of random numbers.
Definition: parallel_rng_types.F:51