tch_plus/vision/
squeezenet.rs

1//! SqueezeNet implementation.
2
3use crate::{nn, nn::Module, nn::ModuleT, Tensor};
4
5fn max_pool2d(xs: &Tensor) -> Tensor {
6    xs.max_pool2d([3, 3], [2, 2], [0, 0], [1, 1], true)
7}
8
9fn fire(p: nn::Path, c_in: i64, c_squeeze: i64, c_exp1: i64, c_exp3: i64) -> impl Module {
10    let cfg3 = nn::ConvConfig { padding: 1, ..Default::default() };
11    let squeeze = nn::conv2d(&p / "squeeze", c_in, c_squeeze, 1, Default::default());
12    let exp1 = nn::conv2d(&p / "expand1x1", c_squeeze, c_exp1, 1, Default::default());
13    let exp3 = nn::conv2d(&p / "expand3x3", c_squeeze, c_exp3, 3, cfg3);
14    nn::func(move |xs| {
15        let xs = xs.apply(&squeeze).relu();
16        Tensor::cat(&[xs.apply(&exp1).relu(), xs.apply(&exp3).relu()], 1)
17    })
18}
19
20fn squeezenet(p: &nn::Path, v1_0: bool, nclasses: i64) -> impl ModuleT {
21    let f_p = p / "features";
22    let c_p = p / "classifier";
23    let initial_conv_cfg = nn::ConvConfig { stride: 2, ..Default::default() };
24    let final_conv_cfg = nn::ConvConfig { stride: 1, ..Default::default() };
25    let features = if v1_0 {
26        nn::seq_t()
27            .add(nn::conv2d(&f_p / "0", 3, 96, 7, initial_conv_cfg))
28            .add_fn(|xs| xs.relu())
29            .add_fn(max_pool2d)
30            .add(fire(&f_p / "3", 96, 16, 64, 64))
31            .add(fire(&f_p / "4", 128, 16, 64, 64))
32            .add(fire(&f_p / "5", 128, 32, 128, 128))
33            .add_fn(max_pool2d)
34            .add(fire(&f_p / "7", 256, 32, 128, 128))
35            .add(fire(&f_p / "8", 256, 48, 192, 192))
36            .add(fire(&f_p / "9", 384, 48, 192, 192))
37            .add(fire(&f_p / "10", 384, 64, 256, 256))
38            .add_fn(max_pool2d)
39            .add(fire(&f_p / "12", 512, 64, 256, 256))
40    } else {
41        nn::seq_t()
42            .add(nn::conv2d(&f_p / "0", 3, 64, 3, initial_conv_cfg))
43            .add_fn(|xs| xs.relu())
44            .add_fn(max_pool2d)
45            .add(fire(&f_p / "3", 64, 16, 64, 64))
46            .add(fire(&f_p / "4", 128, 16, 64, 64))
47            .add_fn(max_pool2d)
48            .add(fire(&f_p / "6", 128, 32, 128, 128))
49            .add(fire(&f_p / "7", 256, 32, 128, 128))
50            .add_fn(max_pool2d)
51            .add(fire(&f_p / "9", 256, 48, 192, 192))
52            .add(fire(&f_p / "10", 384, 48, 192, 192))
53            .add(fire(&f_p / "11", 384, 64, 256, 256))
54            .add(fire(&f_p / "12", 512, 64, 256, 256))
55    };
56    features
57        .add_fn_t(|xs, train| xs.dropout(0.5, train))
58        .add(nn::conv2d(&c_p / "1", 512, nclasses, 1, final_conv_cfg))
59        .add_fn(|xs| xs.relu().adaptive_avg_pool2d([1, 1]).flat_view())
60}
61
62pub fn v1_0(p: &nn::Path, nclasses: i64) -> impl ModuleT {
63    squeezenet(p, true, nclasses)
64}
65
66pub fn v1_1(p: &nn::Path, nclasses: i64) -> impl ModuleT {
67    squeezenet(p, false, nclasses)
68}
tch_plus/vision/squeezenet.rs

tch_plus/vision/
squeezenet.rs
