pub fn samples3(
x: &[f32],
dims: &[usize],
cfg: &Gpu3Config,
threads: usize,
max: usize,
) -> Vec<Sample>Expand description
Training samples for a field at a PSNR/RMSE target (the coefficients the v3 coder would see at the rate-controlled step).
Examples found in repository?
examples/pretrain3.rs (line 50)
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}