Revert "Revert "Implement 3D and transposed 3D convolutions. (#1945)""

This reverts commit b8b47ea.

backend-comparison/Cargo.toml

@@ -72,10 +72,19 @@ name = "conv-transpose2d"
 path = "benches/conv_transpose2d.rs"
 harness = false
 
+[[bench]]
+name = "conv-transpose3d"
+path = "benches/conv_transpose3d.rs"
+harness = false
+
 [[bench]]
 name = "conv2d"
 harness = false
 
+[[bench]]
+name = "conv3d"
+harness = false
+
 [[bench]]
 name = "matmul"
 harness = false

backend-comparison/benches/conv3d.rs

@@ -0,0 +1,109 @@
+use backend_comparison::persistence::save;
+use burn::tensor::{
+ backend::Backend, module::conv3d, ops::ConvOptions, Distribution, Shape, Tensor,
+};
+use burn_common::{
+ benchmark::{run_benchmark, Benchmark},
+ sync_type::SyncType,
+};
+
+pub struct Conv3dBenchmark<B: Backend> {
+ input_shape: Shape<5>,
+ weight_shape: Shape<5>,
+ bias_shape: Shape<1>,
+ options: ConvOptions<3>,
+ device: B::Device,
+}
+
+impl<B: Backend> Benchmark for Conv3dBenchmark<B> {
+ type Args = (Tensor<B, 5>, Tensor<B, 5>, Tensor<B, 1>);
+
+ fn name(&self) -> String {
+ "conv3d".into()
+ }
+
+ fn shapes(&self) -> Vec<Vec<usize>> {
+ vec![
+ self.input_shape.dims.into(),
+ self.weight_shape.dims.into(),
+ self.bias_shape.dims.into(),
+ ]
+ }
+
+ fn execute(&self, (x, w, b): Self::Args) {
+ conv3d(x, w, Some(b), self.options.clone());
+ }
+
+ fn prepare(&self) -> Self::Args {
+ (
+ Tensor::random(
+ self.input_shape.clone(),
+ Distribution::Default,
+ &self.device,
+ ),
+ Tensor::random(
+ self.weight_shape.clone(),
+ Distribution::Default,
+ &self.device,
+ ),
+ Tensor::random(self.bias_shape.clone(), Distribution::Default, &self.device),
+ )
+ }
+
+ fn sync(&self) {
+ B::sync(&self.device, SyncType::Wait)
+ }
+}
+
+#[allow(dead_code)]
+fn bench<B: Backend>(
+ device: &B::Device,
+ feature_name: &str,
+ url: Option<&str>,
+ token: Option<&str>,
+) {
+ // Shapes
+ let batch_size = 16;
+ let channels_in = 16;
+ let channels_out = 16;
+ let depth_in = 16;
+ let height_in = 128;
+ let width_in = 128;
+ let kernel_size_0 = 3;
+ let kernel_size_1 = 3;
+ let kernel_size_2 = 3;
+
+ // Options
+ let strides = [1, 1, 1];
+ let padding = [0, 0, 0];
+ let dilations = [1, 1, 1];
+ let groups = 1;
+ let options = ConvOptions::new(strides, padding, dilations, groups);
+ let benchmark = Conv3dBenchmark::<B> {
+ input_shape: [batch_size, channels_in, depth_in, height_in, width_in].into(),
+ weight_shape: [
+ channels_in,
+ channels_out / groups,
+ kernel_size_0,
+ kernel_size_1,
+ kernel_size_2,
+ ]
+ .into(),
+ bias_shape: [channels_out].into(),
+ options,
+ device: device.clone(),
+ };
+
+ save::<B>(
+ vec![run_benchmark(benchmark)],
+ device,
+ feature_name,
+ url,
+ token,
+ )
+ .unwrap();
+}
+
+fn main() {
+ backend_comparison::bench_on_backend!();
+}

backend-comparison/benches/conv_transpose3d.rs

@@ -0,0 +1,111 @@
+use backend_comparison::persistence::save;
+use burn::tensor::{
+ backend::Backend, module::conv_transpose3d, ops::ConvTransposeOptions, Distribution, Shape,
+ Tensor,
+};
+use burn_common::{
+ benchmark::{run_benchmark, Benchmark},
+ sync_type::SyncType,
+};
+
+pub struct ConvTranspose3dBenchmark<B: Backend> {
+ input_shape: Shape<5>,
+ weight_shape: Shape<5>,
+ bias_shape: Shape<1>,
+ options: ConvTransposeOptions<3>,
+ device: B::Device,
+}
+
+impl<B: Backend> Benchmark for ConvTranspose3dBenchmark<B> {
+ type Args = (Tensor<B, 5>, Tensor<B, 5>, Tensor<B, 1>);
+
+ fn name(&self) -> String {
+ "conv_transpose3d".into()
+ }
+
+ fn shapes(&self) -> Vec<Vec<usize>> {
+ vec![
+ self.input_shape.dims.into(),
+ self.weight_shape.dims.into(),
+ self.bias_shape.dims.into(),
+ ]
+ }
+
+ fn execute(&self, (x, w, b): Self::Args) {
+ conv_transpose3d(x, w, Some(b), self.options.clone());
+ }
+
+ fn prepare(&self) -> Self::Args {
+ (
+ Tensor::random(
+ self.input_shape.clone(),
+ Distribution::Default,
+ &self.device,
+ ),
+ Tensor::random(
+ self.weight_shape.clone(),
+ Distribution::Default,
+ &self.device,
+ ),
+ Tensor::random(self.bias_shape.clone(), Distribution::Default, &self.device),
+ )
+ }
+
+ fn sync(&self) {
+ B::sync(&self.device, SyncType::Wait)
+ }
+}
+
+#[allow(dead_code)]
+fn bench<B: Backend>(
+ device: &B::Device,
+ feature_name: &str,
+ url: Option<&str>,
+ token: Option<&str>,
+) {
+ // Shapes
+ let batch_size = 16;
+ let channels_in = 16;
+ let channels_out = 16;
+ let depth_in = 4;
+ let height_in = 16;
+ let width_in = 16;
+ let kernel_size_0 = 8;
+ let kernel_size_1 = 8;
+ let kernel_size_2 = 8;
+
+ // Options
+ let strides = [1, 1, 1];
+ let padding = [0, 0, 0];
+ let padding_out = [0, 0, 0];
+ let dilations = [1, 1, 1];
+ let groups = 1;
+ let options = ConvTransposeOptions::new(strides, padding, padding_out, dilations, groups);
+ let benchmark = ConvTranspose3dBenchmark::<B> {
+ input_shape: [batch_size, channels_in, depth_in, height_in, width_in].into(),
+ weight_shape: [
+ channels_in,
+ channels_out / groups,
+ kernel_size_0,
+ kernel_size_1,
+ kernel_size_2,
+ ]
+ .into(),
+ bias_shape: [channels_out].into(),
+ options,
+ device: device.clone(),
+ };
+
+ save::<B>(
+ vec![run_benchmark(benchmark)],
+ device,
+ feature_name,
+ url,
+ token,
+ )
+ .unwrap();
+}
+
+fn main() {
+ backend_comparison::bench_on_backend!();
+}

backend-comparison/src/burnbenchapp/base.rs

@@ -104,8 +104,12 @@ enum BenchmarkValues {
  Autodiff,
  #[strum(to_string = "conv-transpose2d")]
  ConvTranspose2d,
+ #[strum(to_string = "conv-transpose3d")]
+ ConvTranspose3d,
  #[strum(to_string = "conv2d")]
  Conv2d,
+ #[strum(to_string = "conv3d")]
+ Conv3d,
 }
 
 pub fn execute() {

burn-book/src/building-blocks/module.md

@@ -183,8 +183,10 @@ Burn comes with built-in modules that you can use to build your own modules.
 | ----------------- | -------------------- |
 | `Conv1d` | `nn.Conv1d` |
 | `Conv2d` | `nn.Conv2d` |
+| `Conv3d` | `nn.Conv3d` |
 | `ConvTranspose1d` | `nn.ConvTranspose1d` |
 | `ConvTranspose2d` | `nn.ConvTranspose2d` |
+| `ConvTranspose3d` | `nn.ConvTranspose3d` |
 
 ### Pooling