committed by
FilippoOlivo
parent
3684782fb5
commit
578c5bc2f4
68
tutorials/tutorial11/tutorial.py
vendored
68
tutorials/tutorial11/tutorial.py
vendored
@@ -1,16 +1,16 @@
|
||||
#!/usr/bin/env python
|
||||
# coding: utf-8
|
||||
|
||||
# # Tutorial: PINA and PyTorch Lightning, training tips and visualizations
|
||||
#
|
||||
# # Tutorial: PINA and PyTorch Lightning, training tips and visualizations
|
||||
#
|
||||
# [](https://colab.research.google.com/github/mathLab/PINA/blob/master/tutorials/tutorial11/tutorial.ipynb)
|
||||
#
|
||||
# In this tutorial, we will delve deeper into the functionality of the `Trainer` class, which serves as the cornerstone for training **PINA** [Solvers](https://mathlab.github.io/PINA/_rst/_code.html#solvers).
|
||||
#
|
||||
#
|
||||
# In this tutorial, we will delve deeper into the functionality of the `Trainer` class, which serves as the cornerstone for training **PINA** [Solvers](https://mathlab.github.io/PINA/_rst/_code.html#solvers).
|
||||
#
|
||||
# The `Trainer` class offers a plethora of features aimed at improving model accuracy, reducing training time and memory usage, facilitating logging visualization, and more thanks to the amazing job done by the PyTorch Lightning team!
|
||||
#
|
||||
#
|
||||
# Our leading example will revolve around solving the `SimpleODE` problem, as outlined in the [*Introduction to PINA for Physics Informed Neural Networks training*](https://github.com/mathLab/PINA/blob/master/tutorials/tutorial1/tutorial.ipynb). If you haven't already explored it, we highly recommend doing so before diving into this tutorial.
|
||||
#
|
||||
#
|
||||
# Let's start by importing useful modules, define the `SimpleODE` problem and the `PINN` solver.
|
||||
|
||||
# In[ ]:
|
||||
@@ -105,7 +105,7 @@ trainer = Trainer(solver=pinn)
|
||||
|
||||
|
||||
# ## Trainer Accelerator
|
||||
#
|
||||
#
|
||||
# When creating the trainer, **by defualt** the `Trainer` will choose the most performing `accelerator` for training which is available in your system, ranked as follow:
|
||||
# 1. [TPU](https://cloud.google.com/tpu/docs/intro-to-tpu)
|
||||
# 2. [IPU](https://www.graphcore.ai/products/ipu)
|
||||
@@ -114,7 +114,7 @@ trainer = Trainer(solver=pinn)
|
||||
# 5. CPU
|
||||
|
||||
# For setting manually the `accelerator` run:
|
||||
#
|
||||
#
|
||||
# * `accelerator = {'gpu', 'cpu', 'hpu', 'mps', 'cpu', 'ipu'}` sets the accelerator to a specific one
|
||||
|
||||
# In[4]:
|
||||
@@ -126,13 +126,13 @@ trainer = Trainer(solver=pinn, accelerator="cpu")
|
||||
# as you can see, even if in the used system `GPU` is available, it is not used since we set `accelerator='cpu'`.
|
||||
|
||||
# ## Trainer Logging
|
||||
#
|
||||
#
|
||||
# In **PINA** you can log metrics in different ways. The simplest approach is to use the `MetricTraker` class from `pina.callbacks` as seen in the [*Introduction to PINA for Physics Informed Neural Networks training*](https://github.com/mathLab/PINA/blob/master/tutorials/tutorial1/tutorial.ipynb) tutorial.
|
||||
#
|
||||
#
|
||||
# However, expecially when we need to train multiple times to get an average of the loss across multiple runs, `pytorch_lightning.loggers` might be useful. Here we will use `TensorBoardLogger` (more on [logging](https://lightning.ai/docs/pytorch/stable/extensions/logging.html) here), but you can choose the one you prefer (or make your own one).
|
||||
#
|
||||
#
|
||||
# We will now import `TensorBoardLogger`, do three runs of training and then visualize the results. Notice we set `enable_model_summary=False` to avoid model summary specifications (e.g. number of parameters), set it to true if needed.
|
||||
#
|
||||
#
|
||||
|
||||
# In[5]:
|
||||
|
||||
@@ -173,17 +173,17 @@ for _ in range(3):
|
||||
|
||||
# Whenever we need to access certain steps of the training for logging, do static modifications (i.e. not changing the `Solver`) or updating `Problem` hyperparameters (static variables), we can use `Callabacks`. Notice that `Callbacks` allow you to add arbitrary self-contained programs to your training. At specific points during the flow of execution (hooks), the Callback interface allows you to design programs that encapsulate a full set of functionality. It de-couples functionality that does not need to be in **PINA** `Solver`s.
|
||||
# Lightning has a callback system to execute them when needed. Callbacks should capture NON-ESSENTIAL logic that is NOT required for your lightning module to run.
|
||||
#
|
||||
#
|
||||
# The following are best practices when using/designing callbacks.
|
||||
#
|
||||
#
|
||||
# * Callbacks should be isolated in their functionality.
|
||||
# * Your callback should not rely on the behavior of other callbacks in order to work properly.
|
||||
# * Do not manually call methods from the callback.
|
||||
# * Directly calling methods (eg. on_validation_end) is strongly discouraged.
|
||||
# * Whenever possible, your callbacks should not depend on the order in which they are executed.
|
||||
#
|
||||
#
|
||||
# We will try now to implement a naive version of `MetricTraker` to show how callbacks work. Notice that this is a very easy application of callbacks, fortunately in **PINA** we already provide more advanced callbacks in `pina.callbacks`.
|
||||
#
|
||||
#
|
||||
# <!-- Suppose we want to log the accuracy on some validation poit -->
|
||||
|
||||
# In[6]:
|
||||
@@ -207,7 +207,7 @@ class NaiveMetricTracker(Callback):
|
||||
)
|
||||
|
||||
|
||||
# Let's see the results when applyed to the `SimpleODE` problem. You can define callbacks when initializing the `Trainer` by the `callbacks` argument, which expects a list of callbacks.
|
||||
# Let's see the results when applyed to the `SimpleODE` problem. You can define callbacks when initializing the `Trainer` by the `callbacks` argument, which expects a list of callbacks.
|
||||
|
||||
# In[7]:
|
||||
|
||||
@@ -240,8 +240,8 @@ trainer.train()
|
||||
trainer.callbacks[0].saved_metrics[:3] # only the first three epochs
|
||||
|
||||
|
||||
# PyTorch Lightning also has some built in `Callbacks` which can be used in **PINA**, [here an extensive list](https://lightning.ai/docs/pytorch/stable/extensions/callbacks.html#built-in-callbacks).
|
||||
#
|
||||
# PyTorch Lightning also has some built in `Callbacks` which can be used in **PINA**, [here an extensive list](https://lightning.ai/docs/pytorch/stable/extensions/callbacks.html#built-in-callbacks).
|
||||
#
|
||||
# We can for example try the `EarlyStopping` routine, which automatically stops the training when a specific metric converged (here the `train_loss`). In order to let the training keep going forever set `max_epochs=-1`.
|
||||
|
||||
# In[ ]:
|
||||
@@ -271,18 +271,18 @@ trainer.train()
|
||||
# As we can see the model automatically stop when the logging metric stopped improving!
|
||||
|
||||
# ## Trainer Tips to Boost Accuracy, Save Memory and Speed Up Training
|
||||
#
|
||||
#
|
||||
# Untill now we have seen how to choose the right `accelerator`, how to log and visualize the results, and how to interface with the program in order to add specific parts of code at specific points by `callbacks`.
|
||||
# Now, we well focus on how boost your training by saving memory and speeding it up, while mantaining the same or even better degree of accuracy!
|
||||
#
|
||||
#
|
||||
#
|
||||
#
|
||||
# There are several built in methods developed in PyTorch Lightning which can be applied straight forward in **PINA**, here we report some:
|
||||
#
|
||||
#
|
||||
# * [Stochastic Weight Averaging](https://pytorch.org/blog/pytorch-1.6-now-includes-stochastic-weight-averaging/) to boost accuracy
|
||||
# * [Gradient Clippling](https://deepgram.com/ai-glossary/gradient-clipping) to reduce computational time (and improve accuracy)
|
||||
# * [Gradient Accumulation](https://lightning.ai/docs/pytorch/stable/common/optimization.html#id3) to save memory consumption
|
||||
# * [Mixed Precision Training](https://lightning.ai/docs/pytorch/stable/common/optimization.html#id3) to save memory consumption
|
||||
#
|
||||
# * [Gradient Accumulation](https://lightning.ai/docs/pytorch/stable/common/optimization.html#id3) to save memory consumption
|
||||
# * [Mixed Precision Training](https://lightning.ai/docs/pytorch/stable/common/optimization.html#id3) to save memory consumption
|
||||
#
|
||||
# We will just demonstrate how to use the first two, and see the results compared to a standard training.
|
||||
# We use the [`Timer`](https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.Timer.html#lightning.pytorch.callbacks.Timer) callback from `pytorch_lightning.callbacks` to take the times. Let's start by training a simple model without any optimization (train for 2000 epochs).
|
||||
|
||||
@@ -347,7 +347,7 @@ print(f'Total training time {trainer.callbacks[0].time_elapsed("train"):.5f} s')
|
||||
# As you can see, the training time does not change at all! Notice that around epoch `1600`
|
||||
# the scheduler is switched from the defalut one `ConstantLR` to the Stochastic Weight Average Learning Rate (`SWALR`).
|
||||
# This is because by default `StochasticWeightAveraging` will be activated after `int(swa_epoch_start * max_epochs)` with `swa_epoch_start=0.7` by default. Finally, the final `mean_loss` is lower when `StochasticWeightAveraging` is used.
|
||||
#
|
||||
#
|
||||
# We will now now do the same but clippling the gradient to be relatively small.
|
||||
|
||||
# In[12]:
|
||||
@@ -376,13 +376,13 @@ print(f'Total training time {trainer.callbacks[0].time_elapsed("train"):.5f} s')
|
||||
|
||||
|
||||
# As we can see we by applying gradient clipping we were able to even obtain lower error!
|
||||
#
|
||||
#
|
||||
# ## What's next?
|
||||
#
|
||||
#
|
||||
# Now you know how to use efficiently the `Trainer` class **PINA**! There are multiple directions you can go now:
|
||||
#
|
||||
# 1. Explore training times on different devices (e.g.) `TPU`
|
||||
#
|
||||
#
|
||||
# 1. Explore training times on different devices (e.g.) `TPU`
|
||||
#
|
||||
# 2. Try to reduce memory cost by mixed precision training and gradient accumulation (especially useful when training Neural Operators)
|
||||
#
|
||||
#
|
||||
# 3. Benchmark `Trainer` speed for different precisions.
|
||||
|
||||
Reference in New Issue
Block a user