(git:374b731)
Loading...
Searching...
No Matches
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! **************************************************************************************************
12MODULE 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
33CONTAINS
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
332END 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
pao_types::pao_env_type
Definition pao_types.F:126
pao_types::training_matrix_type
Definition pao_types.F:44
parallel_rng_types::rng_stream_type
Definition parallel_rng_types.F:157