search

graphnet-team / graphnet

A Deep learning library for neutrino telescopes
https://graphnet-team.github.io/graphnet/
Apache License 2.0
85 stars 86 forks source link

define SWA and EMA model #661

Closed ArturoLlorente closed 4 months ago

ArturoLlorente commented 5 months ago

Stochastic Weight Averaging (SWA) is a regularization technique in PyTorch that improves the generalization performance of deep neural networks by computing the average of multiple weights characterized by different points of the optimization trajectory. It is better explained in the article: PyTorch 1.6 now includes Stochastic Weight Averaging.

In this model, it is also possible to define EMA. Exponential Moving Average (EMA) is a widely known technique to reduce the training time by reducing the number of weight updates needed.

Implementation was made following the example: SWA and EMA pytorch

RasmusOrsoe commented 4 months ago

@ArturoLlorente could you update this PR with the latest changes from main? Then I'll review