pub fn train(samples: &[Sample], cfg: &TrainConfig) -> NetExpand description
Train a network on samples (deterministic for a given seed); returns the quantized integer network.
Examples found in repository?
examples/pretrain3.rs (line 65)
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}More examples
examples/pretrain.rs (line 79)
16fn main() {
17 let a: Vec<String> = std::env::args().skip(1).collect();
18 let out = &a[0];
19 let mut dims = vec![64, 128, 128];
20 let mut psnrs = vec![40.0, 60.0, 80.0, 100.0];
21 let mut cfg = TrainConfig::pretrain();
22 let mut per = 150_000usize;
23 let mut files = Vec::new();
24 let mut i = 1;
25 while i < a.len() {
26 match a[i].as_str() {
27 "--dims" => dims = a[i + 1].split('x').map(|s| s.parse().unwrap()).collect(),
28 "--psnr" => psnrs = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
29 "--hidden" => cfg.hidden = a[i + 1].split(',').map(|s| s.parse().unwrap()).collect(),
30 "--per" => per = a[i + 1].parse().unwrap(),
31 "--epochs" => cfg.epochs = a[i + 1].parse().unwrap(),
32 "--per-epoch" => cfg.per_epoch = a[i + 1].parse().unwrap(),
33 f => {
34 files.push(f.to_string());
35 i += 1;
36 continue;
37 }
38 }
39 i += 2;
40 }
41 let mut samples: Vec<Sample> = Vec::new();
42 let mut seed = 12345u64;
43 for f in &files {
44 let b = std::fs::read(f).unwrap();
45 let x: Vec<f32> = b
46 .as_chunks::<4>()
47 .0
48 .iter()
49 .map(|&c| f32::from_le_bytes(c))
50 .collect();
51 let xd: Vec<f64> = x.iter().map(|&v| v as f64).collect();
52 let co = Coeffs::new(&xd, &dims);
53 for &p in &psnrs {
54 let wc = WavConfig::new(ErrorBound::Psnr(p));
55 let (_, st) = encode_f32(&x, &dims, &wc);
56 let (q, _) = quantize(&co, st.delta, wc.rho);
57 let preds = fit_predictors(&co.p, &q);
58 let mut s = collect(&co.p, &q, &preds);
59 // seeded partial Fisher–Yates subsample
60 let k = per.min(s.len());
61 for j in 0..k {
62 seed = seed
63 .wrapping_mul(6364136223846793005)
64 .wrapping_add(1442695040888963407);
65 let r = j + (seed >> 11) as usize % (s.len() - j);
66 s.swap(j, r);
67 }
68 s.truncate(k);
69 eprintln!(
70 "{f} psnr {p}: delta {:.4e}, hand {:.4} bpv, {} samples",
71 st.delta,
72 st.bits_per_value,
73 s.len()
74 );
75 samples.extend(s);
76 }
77 }
78 let t = std::time::Instant::now();
79 let net = train(&samples, &cfg);
80 eprintln!(
81 "trained {} params on {} samples in {:.1}s: {:.4} bits/sample (float view)",
82 net.num_params(),
83 samples.len(),
84 t.elapsed().as_secs_f64(),
85 mean_bits_float(&net, &samples)
86 );
87 std::fs::write(out, Arc::new(net).to_bytes()).unwrap();
88 let _ = Net::from_bytes(&std::fs::read(out).unwrap()).unwrap();
89}