1use burn::module::{Module, Param};
2use burn::prelude::Backend;
3use burn::tensor::Tensor;
4use burn::tensor::module::{conv_transpose1d, conv1d};
5use burn::tensor::ops::{ConvOptions, ConvTransposeOptions, PadMode};
6
7#[derive(Module, Debug)]
8pub struct WNConv1d<B: Backend> {
9 pub bias: Param<Tensor<B, 1>>,
10
11 pub weight_g: Param<Tensor<B, 3>>,
12 pub weight_v: Param<Tensor<B, 3>>,
13
14 pub stride: usize,
15 pub padding: usize,
16 pub dilation: usize,
17 pub groups: usize,
18 pub causal: bool,
19}
20
21impl<B: Backend> WNConv1d<B> {
22 #[allow(clippy::too_many_arguments)]
23 pub fn new(
24 device: &B::Device,
25 in_channels: usize,
26 out_channels: usize,
27 kernel_size: usize,
28 stride: usize,
29 padding: usize,
30 dilation: usize,
31 groups: usize,
32 causal: bool,
33 ) -> Self {
34 let in_channels_per_group = in_channels / groups;
35
36 Self {
37 bias: Param::from_tensor(Tensor::zeros([out_channels], device)),
38 weight_g: Param::from_tensor(Tensor::ones([out_channels, 1, 1], device)),
39 weight_v: Param::from_tensor(Tensor::zeros(
40 [out_channels, in_channels_per_group, kernel_size],
41 device,
42 )),
43 stride,
44 padding,
45 dilation,
46 groups,
47 causal,
48 }
49 }
50
51 fn compute_weight(&self) -> Tensor<B, 3> {
52 let g = self.weight_g.val();
53 let v = self.weight_v.val();
54
55 let v_norm_sq = v.clone().powf_scalar(2.0).sum_dim(2).sum_dim(1);
56 let v_norm = v_norm_sq.sqrt();
57
58 let out_ch = v_norm.dims()[0];
59 let v_norm = v_norm.reshape([out_ch, 1, 1]);
60
61 g * v / (v_norm + 1e-12)
62 }
63
64 pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
65 let weight = self.compute_weight();
66 let bias = self.bias.val();
67
68 let left_padding = if self.causal {
69 self.dilation * (self.kernel_size().saturating_sub(1))
70 } else {
71 self.padding
72 };
73 let x = if self.causal && left_padding > 0 {
74 x.pad((left_padding, 0, 0, 0), PadMode::Constant(0.0))
75 } else {
76 x
77 };
78 let options = ConvOptions::new(
79 [self.stride],
80 [if self.causal { 0 } else { self.padding }],
81 [self.dilation],
82 self.groups,
83 );
84
85 conv1d(x, weight, Some(bias), options)
86 }
87
88 pub fn kernel_size(&self) -> usize {
89 self.weight_v.dims()[2]
90 }
91}
92
93#[derive(Module, Debug)]
94pub struct WNConvTranspose1d<B: Backend> {
95 pub bias: Param<Tensor<B, 1>>,
96
97 pub weight_g: Param<Tensor<B, 3>>,
98 pub weight_v: Param<Tensor<B, 3>>,
99
100 pub stride: usize,
101 pub padding: usize,
102 pub output_padding: usize,
103 pub dilation: usize,
104 pub groups: usize,
105 pub causal: bool,
106}
107
108impl<B: Backend> WNConvTranspose1d<B> {
109 #[allow(clippy::too_many_arguments)]
110 pub fn new(
111 device: &B::Device,
112 in_channels: usize,
113 out_channels: usize,
114 kernel_size: usize,
115 stride: usize,
116 padding: usize,
117 output_padding: usize,
118 dilation: usize,
119 groups: usize,
120 causal: bool,
121 ) -> Self {
122 Self {
123 bias: Param::from_tensor(Tensor::zeros([out_channels], device)),
124 weight_g: Param::from_tensor(Tensor::ones([in_channels, 1, 1], device)),
125 weight_v: Param::from_tensor(Tensor::zeros(
126 [in_channels, out_channels / groups, kernel_size],
127 device,
128 )),
129 stride,
130 padding,
131 output_padding,
132 dilation,
133 groups,
134 causal,
135 }
136 }
137
138 fn compute_weight(&self) -> Tensor<B, 3> {
139 let g = self.weight_g.val();
140 let v = self.weight_v.val();
141
142 let v_norm_sq = v.clone().powf_scalar(2.0).sum_dim(2).sum_dim(1);
143 let v_norm = v_norm_sq.sqrt();
144
145 let in_ch = v_norm.dims()[0];
146 let v_norm = v_norm.reshape([in_ch, 1, 1]);
147
148 g * v / (v_norm + 1e-12)
149 }
150
151 pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
152 let weight = self.compute_weight();
153 let bias = self.bias.val();
154
155 let options = ConvTransposeOptions::new(
156 [self.stride],
157 [if self.causal { 0 } else { self.padding }],
158 [self.output_padding],
159 [self.dilation],
160 self.groups,
161 );
162
163 let x = conv_transpose1d(x, weight, Some(bias), options);
164 if self.causal {
165 let [batch, channels, time] = x.dims();
166 let crop = self.stride.min(time);
167 x.slice([0..batch, 0..channels, 0..time.saturating_sub(crop)])
168 } else {
169 x
170 }
171 }
172}
173
174pub struct WNConv1dLoadArgs {
175 pub in_channels: usize,
176 pub out_channels: usize,
177 pub kernel_size: usize,
178 pub dilation: usize,
179 pub causal: bool,
180 pub weight_g: (Vec<f32>, Vec<usize>),
181 pub weight_v: (Vec<f32>, Vec<usize>),
182 pub bias: Option<(Vec<f32>, Vec<usize>)>,
183}
184
185pub fn load_wnconv_from_tensors<B: Backend>(
186 device: &B::Device,
187 args: WNConv1dLoadArgs,
188) -> anyhow::Result<WNConv1d<B>> {
189 use burn::tensor::TensorData;
190
191 let (g_data, g_shape) = args.weight_g;
192 let weight_g_tensor = Tensor::<B, 3>::from_data(
193 TensorData::new(g_data, [g_shape[0], g_shape[1], g_shape[2]]),
194 device,
195 );
196
197 let (v_data, v_shape) = args.weight_v;
198 let weight_v_tensor = Tensor::<B, 3>::from_data(
199 TensorData::new(v_data, [v_shape[0], v_shape[1], v_shape[2]]),
200 device,
201 );
202
203 let bias_tensor = if let Some((b_data, b_shape)) = args.bias {
204 Tensor::<B, 1>::from_data(TensorData::new(b_data, [b_shape[0]]), device)
205 } else {
206 Tensor::zeros([args.out_channels], device)
207 };
208
209 Ok(WNConv1d {
210 bias: Param::from_tensor(bias_tensor),
211 weight_g: Param::from_tensor(weight_g_tensor),
212 weight_v: Param::from_tensor(weight_v_tensor),
213 stride: 1,
214 padding: (args.kernel_size / 2) * args.dilation,
215 dilation: args.dilation,
216 groups: 1,
217 causal: args.causal,
218 })
219}
220
221pub struct WNConvTranspose1dLoadArgs {
222 pub out_channels: usize,
223 pub kernel_size: usize,
224 pub stride: usize,
225 pub causal: bool,
226 pub weight_g: (Vec<f32>, Vec<usize>),
227 pub weight_v: (Vec<f32>, Vec<usize>),
228 pub bias: Option<(Vec<f32>, Vec<usize>)>,
229}
230
231pub fn load_wnconv_transpose_from_tensors<B: Backend>(
232 device: &B::Device,
233 args: WNConvTranspose1dLoadArgs,
234) -> anyhow::Result<WNConvTranspose1d<B>> {
235 use burn::tensor::TensorData;
236
237 let (g_data, g_shape) = args.weight_g;
238 let weight_g_tensor = Tensor::<B, 3>::from_data(
239 TensorData::new(g_data, [g_shape[0], g_shape[1], g_shape[2]]),
240 device,
241 );
242
243 let (v_data, v_shape) = args.weight_v;
244 let weight_v_tensor = Tensor::<B, 3>::from_data(
245 TensorData::new(v_data, [v_shape[0], v_shape[1], v_shape[2]]),
246 device,
247 );
248
249 let bias_tensor = if let Some((b_data, b_shape)) = args.bias {
250 Tensor::<B, 1>::from_data(TensorData::new(b_data, [b_shape[0]]), device)
251 } else {
252 Tensor::zeros([args.out_channels], device)
253 };
254
255 Ok(WNConvTranspose1d {
256 bias: Param::from_tensor(bias_tensor),
257 weight_g: Param::from_tensor(weight_g_tensor),
258 weight_v: Param::from_tensor(weight_v_tensor),
259 stride: args.stride,
260 padding: args.kernel_size / 2,
261 output_padding: 0,
262 dilation: 1,
263 groups: 1,
264 causal: args.causal,
265 })
266}
267
268pub fn load_prelu_from_tensor<B: Backend>(
269 device: &B::Device,
270 data: Vec<f32>,
271 shape: Vec<usize>,
272) -> anyhow::Result<super::PReLU<B>> {
273 use burn::module::Param;
274 use burn::tensor::TensorData;
275
276 let tensor = Tensor::<B, 1>::from_data(TensorData::new(data, [shape[0]]), device);
277
278 Ok(super::PReLU {
279 weight: Param::from_tensor(tensor),
280 })
281}
282
283#[derive(Module, Debug)]
284pub struct PlainConv1d<B: Backend> {
285 pub weight: Param<Tensor<B, 3>>,
286 pub bias: Param<Tensor<B, 1>>,
287
288 pub stride: usize,
289 pub padding: usize,
290 pub dilation: usize,
291 pub groups: usize,
292 pub causal: bool,
293}
294
295#[derive(Module, Debug)]
296pub struct PostProcessor<B: Backend> {
297 pub conv: PlainConv1d<B>,
298 pub activation: super::PReLU<B>,
299 pub num_samples: usize,
300}
301
302impl<B: Backend> PostProcessor<B> {
303 pub fn new(device: &B::Device, channels: usize, num_samples: usize) -> Self {
304 Self {
305 conv: PlainConv1d::new(device, channels, channels, 7, 1, 3, 1, 1, true),
306 activation: super::PReLU::new(device),
307 num_samples,
308 }
309 }
310
311 pub fn load_from_tensors<F>(
312 device: &B::Device,
313 get_tensor: &F,
314 prefix: &str,
315 channels: usize,
316 num_samples: usize,
317 ) -> anyhow::Result<Self>
318 where
319 F: Fn(&str) -> Option<(Vec<f32>, Vec<usize>)>,
320 {
321 let conv = if let Some((w_data, w_shape)) = get_tensor(&format!("{}.conv.weight", prefix)) {
322 let b = get_tensor(&format!("{}.conv.bias", prefix));
323 load_conv1d_from_tensors(
324 device,
325 Conv1dLoadArgs {
326 in_channels: channels,
327 out_channels: channels,
328 kernel_size: 7,
329 causal: true,
330 weight: (w_data, w_shape),
331 bias: b,
332 },
333 )?
334 } else {
335 PlainConv1d::new(device, channels, channels, 7, 1, 3, 1, 1, true)
336 };
337 let activation =
338 if let Some((data, shape)) = get_tensor(&format!("{}.activation.weight", prefix)) {
339 load_prelu_from_tensor(device, data, shape)?
340 } else {
341 super::PReLU::new(device)
342 };
343 Ok(Self {
344 conv,
345 activation,
346 num_samples,
347 })
348 }
349
350 pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
351 let [batch, channels, time] = x.dims();
352 let x = x.swap_dims(1, 2);
353 let x = if self.num_samples <= 1 {
354 x
355 } else {
356 let repeated: Tensor<B, 4> = x.unsqueeze_dim::<4>(2).repeat_dim(2, self.num_samples);
357 repeated.reshape([batch, time * self.num_samples, channels])
358 };
359 let x = x.swap_dims(1, 2);
360 let x = self.conv.forward(x);
361 self.activation.forward(x)
362 }
363}
364
365impl<B: Backend> PlainConv1d<B> {
366 #[allow(clippy::too_many_arguments)]
367 pub fn new(
368 device: &B::Device,
369 in_channels: usize,
370 out_channels: usize,
371 kernel_size: usize,
372 stride: usize,
373 padding: usize,
374 dilation: usize,
375 groups: usize,
376 causal: bool,
377 ) -> Self {
378 let in_channels_per_group = in_channels / groups;
379
380 Self {
381 weight: Param::from_tensor(Tensor::zeros(
382 [out_channels, in_channels_per_group, kernel_size],
383 device,
384 )),
385 bias: Param::from_tensor(Tensor::zeros([out_channels], device)),
386 stride,
387 padding,
388 dilation,
389 groups,
390 causal,
391 }
392 }
393
394 pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
395 let weight = self.weight.val();
396 let bias = self.bias.val();
397 let left_padding = if self.causal {
398 self.dilation * (self.kernel_size().saturating_sub(1))
399 } else {
400 self.padding
401 };
402 let x = if self.causal && left_padding > 0 {
403 x.pad((left_padding, 0, 0, 0), PadMode::Constant(0.0))
404 } else {
405 x
406 };
407 let options = ConvOptions::new(
408 [self.stride],
409 [if self.causal { 0 } else { self.padding }],
410 [self.dilation],
411 self.groups,
412 );
413 conv1d(x, weight, Some(bias), options)
414 }
415
416 pub fn kernel_size(&self) -> usize {
417 self.weight.dims()[2]
418 }
419}
420
421pub struct Conv1dLoadArgs {
422 pub in_channels: usize,
423 pub out_channels: usize,
424 pub kernel_size: usize,
425 pub causal: bool,
426 pub weight: (Vec<f32>, Vec<usize>),
427 pub bias: Option<(Vec<f32>, Vec<usize>)>,
428}
429
430pub fn load_conv1d_from_tensors<B: Backend>(
431 device: &B::Device,
432 args: Conv1dLoadArgs,
433) -> anyhow::Result<PlainConv1d<B>> {
434 use burn::tensor::TensorData;
435
436 let (w_data, w_shape) = args.weight;
437 let weight_tensor = Tensor::<B, 3>::from_data(
438 TensorData::new(w_data, [w_shape[0], w_shape[1], w_shape[2]]),
439 device,
440 );
441
442 let bias_tensor = if let Some((b_data, b_shape)) = args.bias {
443 Tensor::<B, 1>::from_data(TensorData::new(b_data, [b_shape[0]]), device)
444 } else {
445 Tensor::zeros([args.out_channels], device)
446 };
447
448 Ok(PlainConv1d {
449 weight: Param::from_tensor(weight_tensor),
450 bias: Param::from_tensor(bias_tensor),
451 stride: 1,
452 padding: args.kernel_size / 2,
453 dilation: 1,
454 groups: 1,
455 causal: args.causal,
456 })
457}
458
459#[cfg(test)]
460mod tests {
461 #[test]
462 fn transposed_conv_uses_burn_channel_layout() {
463 let device = burn::backend::ndarray::NdArrayDevice::default();
464 let layer = super::WNConvTranspose1d::<burn::backend::ndarray::NdArray<f32>>::new(
465 &device, 128, 2048, 5, 1, 2, 0, 1, 1, false,
466 );
467
468 assert_eq!(layer.weight_g.val().dims(), [128, 1, 1]);
469 assert_eq!(layer.weight_v.val().dims(), [128, 2048, 5]);
470 }
471}