pretrain3/
pretrain3.rs

1//! `pretrain3` — train the TCG3 learned model offline.
2//!
3//! ```text
4//! pretrain3 OUT.tcm --hidden 32,16 --psnr 40,60,80,100 --lanes 8 --block 32 a.f32:64x128x128 b.f32:64x128x128 ...
5//! ```
6use cram_codec::codec::ErrorBound;
7use cram_codec::gpuwav3::{samples3, Gpu3Config};
8use cram_codec::learned::{mean_bits_float, train, TrainConfig};
9
10fn opt(args: &[String], k: &str, d: &str) -> String {
11    args.iter()
12        .position(|a| a == k)
13        .map(|i| args[i + 1].clone())
14        .unwrap_or_else(|| d.to_string())
15}
16
17fn main() {
18    let args: Vec<String> = std::env::args().skip(1).collect();
19    let out = &args[0];
20    let hidden: Vec<usize> = opt(&args, "--hidden", "32,16")
21        .split(',')
22        .map(|v| v.parse().unwrap())
23        .collect();
24    let psnr: Vec<f64> = opt(&args, "--psnr", "40,50,60,70,80,90,100")
25        .split(',')
26        .map(|v| v.parse().unwrap())
27        .collect();
28    let lanes: usize = opt(&args, "--lanes", "8").parse().unwrap();
29    let block: usize = opt(&args, "--block", "32").parse().unwrap();
30    let per: usize = opt(&args, "--per", "100000").parse().unwrap();
31    let epochs: usize = opt(&args, "--epochs", "8").parse().unwrap();
32    let mut smp = Vec::new();
33    for a in args.iter().filter(|a| a.contains(".f32:")) {
34        let (path, d) = a.split_once(':').unwrap();
35        let dims: Vec<usize> = d.split('x').map(|s| s.parse().unwrap()).collect();
36        let b = std::fs::read(path).unwrap();
37        let x: Vec<f32> = b
38            .as_chunks::<4>()
39            .0
40            .iter()
41            .map(|&c| f32::from_le_bytes(c))
42            .collect();
43        for &p in &psnr {
44            let cfg = Gpu3Config {
45                block,
46                block2d: block * 4,
47                lanes,
48                ..Gpu3Config::new(ErrorBound::Psnr(p))
49            };
50            let s = samples3(&x, &dims, &cfg, 16, per);
51            eprintln!("{path} {p} dB: {} samples", s.len());
52            smp.extend(s);
53        }
54    }
55    let tc = TrainConfig {
56        hidden,
57        epochs,
58        per_epoch: 3_000_000,
59        batch: 256,
60        lr: 3e-3,
61        apm: false,
62        seed: 7,
63    };
64    let t = std::time::Instant::now();
65    let net = train(&smp, &tc);
66    eprintln!(
67        "trained on {} samples in {:.1}s: {:.4} bits/sample (float view), {} params",
68        smp.len(),
69        t.elapsed().as_secs_f64(),
70        mean_bits_float(&net, &smp),
71        net.num_params()
72    );
73    std::fs::write(out, net.to_bytes()).unwrap();
74}
pretrain3/pretrain3.rs

pretrain3/
pretrain3.rs
