burn_dragon_train 0.5.0

Training utilities for burn_dragon
Documentation 
1
2
3
4
5
6
7
8
9
10
11

use crate::train::prelude::*;

pub fn adamw_config_from_optimizer(optimizer_cfg: &OptimizerConfig) -> AdamWConfig {
    let mut config = AdamWConfig::new().with_weight_decay(optimizer_cfg.weight_decay);
    if let Some(clip) = optimizer_cfg.grad_clip_norm {
        config = config.with_grad_clipping(Some(GradientClippingConfig::Norm(clip)));
    } else if let Some(clip) = optimizer_cfg.grad_clip_value {
        config = config.with_grad_clipping(Some(GradientClippingConfig::Value(clip)));
    }
    config
}