burn-vision 0.21.0

Vision processing operations for burn tensors
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
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

use burn_core as burn;

use crate::PaddingConfig2d;
use crate::conv::{Conv2d, Conv2dConfig};
use burn::module::Module;
use burn::tensor::{
    Tensor,
    activation::relu,
    backend::Backend,
    module::{avg_pool2d, max_pool2d},
};

/// VGG19 feature extractor for the Gram Matrix Loss.
///
/// This module implements the VGG19 architecture up to the 5th convolutional block.
/// It is specifically tailored for Neural Style Transfer and Gram Matrix Loss,
/// extracting and flattening features from the following 5 layers:
/// - `conv1_1`
/// - `conv2_1`
/// - `conv3_1`
/// - `conv4_1`
/// - `conv5_1`
#[derive(Module, Debug)]
pub struct Vgg19<B: Backend> {
    use_avg_pool: bool,

    // Block 1
    // Field is made public for testing whether the weights are frozen or not
    pub conv1_1: Conv2d<B>,
    conv1_2: Conv2d<B>,

    // Block 2
    conv2_1: Conv2d<B>,
    conv2_2: Conv2d<B>,

    // Block 3
    conv3_1: Conv2d<B>,
    conv3_2: Conv2d<B>,
    conv3_3: Conv2d<B>,
    conv3_4: Conv2d<B>,

    // Block 4
    conv4_1: Conv2d<B>,
    conv4_2: Conv2d<B>,
    conv4_3: Conv2d<B>,
    conv4_4: Conv2d<B>,

    // Block 5
    conv5_1: Conv2d<B>,
}

impl<B: Backend> Vgg19<B> {
    /// Creates a new VGG19 feature extractor.
    ///
    /// The network is initialized with standard VGG19 configurations (3x3 kernels,
    /// stride 1, padding 1). Note that the weights are randomly initialized here so
    /// they should be overwritten by `load_vgg19_weights` before use.
    pub fn new(use_avg_pool: bool, device: &B::Device) -> Self {
        // All convolutions use a kernel size of 3 by 3, stride of 1, and
        // padding of 1.
        // This combination of kernel size and padding preserves input
        // dimensions. Thus, `PaddingConfig2d::Same` can be used instead.
        let conv_config = |in_ch, out_ch| {
            Conv2dConfig::new([in_ch, out_ch], [3, 3])
                .with_stride([1, 1])
                .with_padding(PaddingConfig2d::Same)
                .init(device)
        };

        Self {
            use_avg_pool,
            // Block 1
            conv1_1: conv_config(3, 64),
            conv1_2: conv_config(64, 64),
            // Block 2
            conv2_1: conv_config(64, 128),
            conv2_2: conv_config(128, 128),
            // Block 3
            conv3_1: conv_config(128, 256),
            conv3_2: conv_config(256, 256),
            conv3_3: conv_config(256, 256),
            conv3_4: conv_config(256, 256),
            // Block 4
            conv4_1: conv_config(256, 512),
            conv4_2: conv_config(512, 512),
            conv4_3: conv_config(512, 512),
            conv4_4: conv_config(512, 512),
            // Block 5
            conv5_1: conv_config(512, 512),
        }
    }

    /// Performs a forward pass to extract features for the Gram Matrix Loss.
    ///
    /// # Arguments
    ///
    /// - `x` - Input image tensor of shape `[batch_size, 3, height, width]`.
    ///
    /// # Returns
    ///
    /// A tuple containing:
    /// - `features`: A `Vec` of 5 tensors, each representing the flattened feature map
    ///    from one of the target layers. Shape of each tensor: `[batch_size, channels, height * width]`.
    /// - `normalization_factors`: A `Vec` of 5 `f32` values, representing the normalization
    ///    factor `4 * N^2 * M^2` for each layer, used to scale the Gram matrix loss.
    pub fn forward(&self, x: Tensor<B, 4>) -> Vec<Tensor<B, 3>> {
        let pool_2d = |x| {
            if self.use_avg_pool {
                avg_pool2d(x, [2, 2], [2, 2], [0, 0], false, false)
            } else {
                max_pool2d(x, [2, 2], [2, 2], [0, 0], [1, 1], false)
            }
        };

        let mut features = Vec::with_capacity(5);

        // Block 1
        let x1_1 = relu(self.conv1_1.forward(x));
        let flattened_x1_1 = x1_1.clone().flatten(2, 3);
        features.push(flattened_x1_1);
        let x1_2 = relu(self.conv1_2.forward(x1_1));
        let x1 = pool_2d(x1_2);

        // Block 2
        let x2_1 = relu(self.conv2_1.forward(x1));
        let flattened_x2_1 = x2_1.clone().flatten(2, 3);
        features.push(flattened_x2_1);
        let x2_2 = relu(self.conv2_2.forward(x2_1));
        let x2 = pool_2d(x2_2);

        // Block 3
        let x3_1 = relu(self.conv3_1.forward(x2));
        let flattened_x3_1 = x3_1.clone().flatten(2, 3);
        features.push(flattened_x3_1);
        let x3_2 = relu(self.conv3_2.forward(x3_1));
        let x3_3 = relu(self.conv3_3.forward(x3_2));
        let x3_4 = relu(self.conv3_4.forward(x3_3));
        let x3 = pool_2d(x3_4);

        // Block 4
        let x4_1 = relu(self.conv4_1.forward(x3));
        let flattened_x4_1 = x4_1.clone().flatten(2, 3);
        features.push(flattened_x4_1);
        let x4_2 = relu(self.conv4_2.forward(x4_1));
        let x4_3 = relu(self.conv4_3.forward(x4_2));
        let x4_4 = relu(self.conv4_4.forward(x4_3));
        let x4 = pool_2d(x4_4);

        // Block 5
        let x5_1 = relu(self.conv5_1.forward(x4));
        let flattened_x5_1 = x5_1.flatten(2, 3);
        features.push(flattened_x5_1);

        features
    }
}