burn-train 0.21.0-pre.5

Training crate for the Burn framework
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

//! Pretrained weights loading for DISTS.

use burn_core as burn;

use burn::tensor::backend::Backend;
use burn_std::network::downloader::download_file_as_bytes;
use burn_store::{ModuleSnapshot, PytorchStore};
use std::fs::{File, create_dir_all};
use std::io::Write;
use std::path::PathBuf;

use super::metric::Dists;

/// URL for pretrained DISTS alpha/beta weights from the official repository.
/// Reference: https://github.com/dingkeyan93/DISTS
const DISTS_WEIGHTS_URL: &str =
    "https://github.com/dingkeyan93/DISTS/raw/master/DISTS_pytorch/weights.pt";

/// URL for ImageNet pretrained VGG16 backbone weights from PyTorch.
const VGG16_IMAGENET_URL: &str = "https://download.pytorch.org/models/vgg16-397923af.pth";

/// Get the cache directory for DISTS weights.
fn get_cache_dir() -> PathBuf {
    let cache_dir = dirs::cache_dir()
        .expect("Could not get cache directory")
        .join("burn-dataset")
        .join("dists");

    if !cache_dir.exists() {
        create_dir_all(&cache_dir).expect("Failed to create cache directory");
    }

    cache_dir
}

/// Download file if not cached.
fn download_if_needed(url: &str, cache_path: &PathBuf, message: &str) {
    if !cache_path.exists() {
        let bytes = download_file_as_bytes(url, message);
        let mut file = File::create(cache_path).expect("Failed to create cache file");
        file.write_all(&bytes).expect("Failed to write weights");
    }
}

/// Download and load pretrained weights into a DISTS module.
///
/// This loads both:
/// 1. ImageNet pretrained VGG16 backbone weights
/// 2. DISTS trained alpha/beta weights
///
/// Weights are cached in the user's cache directory to avoid re-downloading.
///
/// # Arguments
///
/// * `dists` - The DISTS module to load weights into.
///
/// # Returns
///
/// The DISTS module with loaded pretrained weights.
pub fn load_pretrained_weights<B: Backend>(mut dists: Dists<B>) -> Dists<B> {
    let cache_dir = get_cache_dir();

    // Step 1: Download and load VGG16 ImageNet backbone weights
    let vgg_cache_path = cache_dir.join("vgg16_backbone.pth");
    download_if_needed(
        VGG16_IMAGENET_URL,
        &vgg_cache_path,
        "Downloading VGG16 ImageNet weights for DISTS...",
    );

    // Step 2: Download DISTS alpha/beta weights
    let dists_cache_path = cache_dir.join("dists_weights.pt");
    download_if_needed(
        DISTS_WEIGHTS_URL,
        &dists_cache_path,
        "Downloading DISTS alpha/beta weights...",
    );

    // Load VGG16 backbone weights first
    dists = load_vgg16_backbone_weights(dists, &vgg_cache_path);

    // Then load DISTS alpha/beta weights
    dists = load_dists_weights(dists, &dists_cache_path);

    dists
}

/// Load VGG16 ImageNet pretrained backbone weights.
fn load_vgg16_backbone_weights<B: Backend>(mut dists: Dists<B>, cache_path: &PathBuf) -> Dists<B> {
    let mut store = PytorchStore::from_file(cache_path)
        .allow_partial(true)
        .skip_enum_variants(true)
        // VGG16 features.X -> extractor.convY_Z
        .with_key_remapping(r"^features\.0\.", "extractor.conv1_1.")
        .with_key_remapping(r"^features\.2\.", "extractor.conv1_2.")
        .with_key_remapping(r"^features\.5\.", "extractor.conv2_1.")
        .with_key_remapping(r"^features\.7\.", "extractor.conv2_2.")
        .with_key_remapping(r"^features\.10\.", "extractor.conv3_1.")
        .with_key_remapping(r"^features\.12\.", "extractor.conv3_2.")
        .with_key_remapping(r"^features\.14\.", "extractor.conv3_3.")
        .with_key_remapping(r"^features\.17\.", "extractor.conv4_1.")
        .with_key_remapping(r"^features\.19\.", "extractor.conv4_2.")
        .with_key_remapping(r"^features\.21\.", "extractor.conv4_3.")
        .with_key_remapping(r"^features\.24\.", "extractor.conv5_1.")
        .with_key_remapping(r"^features\.26\.", "extractor.conv5_2.")
        .with_key_remapping(r"^features\.28\.", "extractor.conv5_3.");

    let result = dists.load_from(&mut store);
    if let Err(e) = result {
        log::warn!("Some VGG16 backbone weights could not be loaded: {:?}", e);
    }

    dists
}

/// Load DISTS trained alpha/beta weights.
fn load_dists_weights<B: Backend>(mut dists: Dists<B>, cache_path: &PathBuf) -> Dists<B> {
    let mut store = PytorchStore::from_file(cache_path)
        .allow_partial(true)
        .skip_enum_variants(true);

    let result = dists.load_from(&mut store);
    if let Err(e) = result {
        log::warn!("Some DISTS weights could not be loaded: {:?}", e);
    }

    dists
}