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 24 : 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 12 : 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 5870400 : 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
|
