Coverage for pySDC/tutorial/step_7/D_pySDC_with_PyTorch.py: 100%
44 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 13:41 +0000
1# ---
2# jupyter:
3# jupytext:
4# formats: py:percent
5# kernelspec:
6# display_name: Python 3
7# name: python3
8# language_info:
9# name: python
10# ---
12# %% [markdown]
13# # Part D: pySDC and PyTorch
14#
15# PyTorch is a library for machine learning. Its data structure is called a tensor, which runs on CPUs as well as
16# GPUs, with access to all of PyTorch's machine learning methods. Since the potential for use in pySDC is very large,
17# we have started on a data type that uses PyTorch tensors throughout pySDC.
18#
19# This example trains a network to predict the results of implicit Euler solves for the heat equation. It is too
20# simple to do anything useful, but demonstrates how to use tensors in pySDC and then apply the enormous PyTorch
21# infrastructure. This is work in progress in very early stages! The tensor data type is the simplest possible
22# implementation, rather than an efficient one. If you want to work on this, your input is appreciated!
23#
24# The tensor data type, the problem class `Heat1DFDTensor` using it and the network `HeatEquationModel` live in the
25# helper module `torch_heat.py` next to this tutorial.
27# %%
28from pathlib import Path
30import numpy as np
31import torch
32import torch.nn as nn
33import torch.optim as optim
34from pySDC.tutorial.step_7.torch_heat import Heat1DFDTensor, HeatEquationModel
36# %% [markdown]
37# ## Training at the collocation nodes
38#
39# For a first proof of concept, we train the model specifically for the collocation nodes we use in SDC. If
40# successful, the initial guess would already be the exact solution, and we would need no SDC iterations.
41#
42# What we find is that we can train the network to predict the solution to one very specific problem rather well.
43# However, if we train for something else, here solving for a different step size, the model only predicts what it
44# was trained for last, and loses the ability to solve for what it learned before. This is solely because we chose
45# an overly simple model that is unsuitable to the task at hand, which is likely easily solved with a bit of
46# patience. So take this as a demonstration of the interface between pySDC and PyTorch, and if you want to do a
47# project with it, as a starting point for things that actually do something!
48#
49# The output shows the training loss during training and, after each of three training sessions, the error of the
50# prediction with the current state of the network. To demonstrate the forgetfulness, we finally print the error of
51# all learned predictions after training is complete.
54# %%
55def train_at_collocation_nodes():
56 """
57 For the first proof of concept, we want to train the model specifically to the collocation nodes we use in SDC.
58 If successful, the initial guess would already be the exact solution and we would need no SDC iterations.
60 What we find is that we can train the network to predict the solution to one very specific problem rather well.
61 See the error during training for what happens when we ask the network to solve for exactly what it just trained.
62 However, if we train for something else, i.e. solving to a different step size in this case, we can only use the
63 model to predict the solution of what it's been trained for last and it loses the ability to solve for previously
64 learned things. This is solely because we chose an overly simple model that is unsuitable to the task at hand and
65 is likely easily solved with a bit of patience. This is just a demonstration of the interface between pySDC and
66 PyTorch. If you want to do a project with this, feel free to take this as a starting point and do things that
67 actually do something!
69 The output shows the training loss during training and, after each of three training sessions is complete, the error
70 of the prediction with the current state of the network. To demonstrate the forgetfulness, we finally print the
71 error of all learned predictions after training is complete.
72 """
73 out = ''
74 errors_mid_training = []
75 errors_post_training = []
77 # instantiate the pySDC problem and a model for PyTorch
78 problem = Heat1DFDTensor()
79 model = HeatEquationModel(problem)
81 # setup neural network
82 lr = 0.001
83 num_epochs = 250
84 criterion = nn.MSELoss()
85 optimizer = optim.Adam(model.parameters(), lr=lr)
87 # setup initial conditions
88 t = 0
89 initial_condition = problem.u_exact(t)
91 # train the model to predict the solution at certain collocation nodes
92 collocation_nodes = np.array([0.15505102572168285, 0.6449489742783183, 1]) * 1e-2
93 for dt in collocation_nodes:
95 # get target condition from implicit Euler step
96 target_condition = problem.solve_system(initial_condition, dt, initial_condition, t)
98 # do the training
99 for epoch in range(num_epochs):
100 predicted_state = model(initial_condition, t, dt)
101 loss = criterion(predicted_state.float(), target_condition.float())
103 optimizer.zero_grad()
104 loss.backward()
105 optimizer.step()
107 if (epoch + 1) % 50 == 0:
108 out += f'Training for {dt=:.2e}: Epoch [{epoch+1:4d}/{num_epochs:4d}], Loss: {loss.item():.4e}\n'
110 # evaluate model to compute error
111 model_prediction = model(initial_condition, t, dt)
112 errors_mid_training += [abs(target_condition - model_prediction)]
113 out += f'Error of prediction at {dt:.2e} during training: {abs(target_condition-model_prediction):.2e}\n'
115 # compare model and problem
116 for dt in collocation_nodes:
117 target_condition = problem.solve_system(initial_condition, dt, initial_condition, t)
118 model_prediction = model(initial_condition, t, dt)
119 errors_post_training += [abs(target_condition - model_prediction)]
120 out += f'Error of prediction at {dt:.2e} after training: {abs(target_condition-model_prediction):.2e}\n'
122 print(out)
123 Path('data').mkdir(parents=True, exist_ok=True)
124 with open('data/step_7_D_out.txt', 'w') as file:
125 file.write(out)
127 # test that the training went as expected
128 assert np.greater([1e-2, 1e-4, 1e-5], errors_mid_training).all(), 'Errors during training are larger than expected'
129 assert np.greater([1e0, 1e0, 1e-5], errors_post_training).all(), 'Errors after training are larger than expected'
131 # save the model to use it throughout pySDC
132 torch.save(model.state_dict(), 'data/heat_equation_model.pth')
135if __name__ == '__main__':
136 train_at_collocation_nodes()
138# %% [markdown]
139# ## Results
140#
141# PyTorch does not run in the browser, nor in the environment this website is built in. These are the results of our
142# CI, which runs this part in an environment with PyTorch, in the run that built this page:
143#
144# :::{literalinclude} /../../data/step_7_D_out.txt
145# :language: text
146# :::