-
Notifications
You must be signed in to change notification settings - Fork 364
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
1 parent
a6e3b4e
commit 5bbc5ea
Showing
12 changed files
with
320 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Binary file not shown.
58 changes: 58 additions & 0 deletions
58
crates/burn-import/onnx-tests/tests/avg_pool1d/avg_pool1d.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,58 @@ | ||
| #!/usr/bin/env python3 | ||
|
|
||
| # used to generate model: avg_pool1d.onnx | ||
|
|
||
| import torch | ||
| import torch.nn as nn | ||
|
|
||
|
|
||
| class Model(nn.Module): | ||
| def __init__(self): | ||
| super(Model, self).__init__() | ||
|
|
||
| self.pool1 = nn.AvgPool1d(4, stride=2) | ||
|
|
||
| self.pool2 = nn.AvgPool1d(4, stride=2, padding=2, count_include_pad=True) | ||
|
|
||
| self.pool3 = nn.AvgPool1d(4, stride=2, padding=2, count_include_pad=False) | ||
|
|
||
| def forward(self, x1, x2, x3): | ||
| y1 = self.pool1(x1) | ||
| y2 = self.pool2(x2) | ||
| y3 = self.pool3(x3) | ||
| return y1, y2, y3 | ||
|
|
||
|
|
||
| def main(): | ||
| # Set seed for reproducibility | ||
| torch.manual_seed(1) | ||
|
|
||
| # Print options | ||
| torch.set_printoptions(precision=3) | ||
|
|
||
| # Export to onnx | ||
| model = Model() | ||
| model.eval() | ||
| device = torch.device("cpu") | ||
|
|
||
| file_name = "avg_pool1d.onnx" | ||
| input1 = torch.randn(1, 5, 5, device=device) | ||
| torch.onnx.export(model, (input1, input1, input1), file_name, | ||
| verbose=False, opset_version=16) | ||
|
|
||
| print("Finished exporting model to {}".format(file_name)) | ||
|
|
||
| # Output some test data for use in the test | ||
| print("Test input data shape: {}".format(input1.shape)) | ||
| print("Test input data: {}".format(input1)) | ||
| output1, output2, output3 = model.forward(input1, input1, input1) | ||
| print("Test output1 data shape: {}".format(output1.shape)) | ||
| print("Test output2 data shape: {}".format(output2.shape)) | ||
| print("Test output3 data shape: {}".format(output3.shape)) | ||
| print("Test output1: {}".format(output1)) | ||
| print("Test output2: {}".format(output2)) | ||
| print("Test output3: {}".format(output3)) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| main() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,155 @@ | ||
| use proc_macro2::TokenStream; | ||
| use quote::quote; | ||
|
|
||
| use burn::{nn::pool::AvgPool1dConfig, record::PrecisionSettings}; | ||
|
|
||
| use super::{Node, NodeCodegen}; | ||
| use crate::burn::{BurnImports, OtherType, Scope, TensorType, ToTokens, Type}; | ||
|
|
||
| #[derive(Debug, Clone)] | ||
| pub struct AvgPool1dNode { | ||
| pub field: OtherType, | ||
| pub input: TensorType, | ||
| pub output: TensorType, | ||
| pub config: AvgPool1dConfig, | ||
| } | ||
|
|
||
| impl AvgPool1dNode { | ||
| pub fn new<S: AsRef<str>>( | ||
| name: S, | ||
| input: TensorType, | ||
| output: TensorType, | ||
| config: AvgPool1dConfig, | ||
| ) -> Self { | ||
| Self { | ||
| field: OtherType::new( | ||
| name, | ||
| quote! { | ||
| AvgPool1d | ||
| }, | ||
| ), | ||
| input, | ||
| output, | ||
| config, | ||
| } | ||
| } | ||
| } | ||
|
|
||
| impl<PS: PrecisionSettings> NodeCodegen<PS> for AvgPool1dNode { | ||
| fn input_types(&self) -> Vec<Type> { | ||
| vec![Type::Tensor(self.input.clone())] | ||
| } | ||
| fn output_types(&self) -> Vec<Type> { | ||
| vec![Type::Tensor(self.output.clone())] | ||
| } | ||
| fn field_type(&self) -> Option<Type> { | ||
| Some(Type::Other(self.field.clone())) | ||
| } | ||
|
|
||
| fn field_init(&self) -> Option<TokenStream> { | ||
| let name = &self.field.name; | ||
| let kernel_size = self.config.kernel_size.to_tokens(); | ||
| let strides = self.config.stride.to_tokens(); | ||
| let padding = self.config.padding.to_tokens(); | ||
| let count_include_pad = self.config.count_include_pad; | ||
|
|
||
| let tokens = quote! { | ||
| let #name = AvgPool1dConfig::new(#kernel_size) | ||
| .with_stride(#strides) | ||
| .with_padding(#padding) | ||
| .with_count_include_pad(#count_include_pad) | ||
| .init(); | ||
| }; | ||
|
|
||
| Some(tokens) | ||
| } | ||
|
|
||
| fn forward(&self, scope: &mut Scope, node_position: usize) -> TokenStream { | ||
| let input = scope.tensor_use_owned(&self.input, node_position); | ||
| let output = &self.output.name; | ||
| let field = &self.field.name; | ||
|
|
||
| quote! { | ||
| let #output = self.#field.forward(#input); | ||
| } | ||
| } | ||
|
|
||
| fn register_imports(&self, imports: &mut BurnImports) { | ||
| imports.register("burn::nn::PaddingConfig1d"); | ||
| imports.register("burn::nn::pool::AvgPool1d"); | ||
| imports.register("burn::nn::pool::AvgPool1dConfig"); | ||
| } | ||
|
|
||
| fn into_node(self) -> Node<PS> { | ||
| Node::AvgPool1d(self) | ||
| } | ||
|
|
||
| fn field_serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> { | ||
| S::serialize_none(serializer) | ||
| } | ||
| } | ||
|
|
||
| #[cfg(test)] | ||
| mod tests { | ||
| use super::*; | ||
| use crate::burn::{graph::BurnGraph, node::test::assert_tokens, TensorType}; | ||
| use burn::{nn::PaddingConfig1d, record::FullPrecisionSettings}; | ||
|
|
||
| #[test] | ||
| fn test_codegen() { | ||
| let mut graph = BurnGraph::<FullPrecisionSettings>::default(); | ||
|
|
||
| graph.register(AvgPool1dNode::new( | ||
| "avg_pool1d", | ||
| TensorType::new_float("input", 3), | ||
| TensorType::new_float("output", 3), | ||
| AvgPool1dConfig::new(3) | ||
| .with_stride(1) | ||
| .with_padding(PaddingConfig1d::Valid), | ||
| )); | ||
|
|
||
| graph.register_input_output(vec!["input".to_string()], vec!["output".to_string()]); | ||
|
|
||
| let expected = quote! { | ||
| use burn::{ | ||
| module::Module, | ||
| tensor::{backend::Backend, Tensor}, | ||
| }; | ||
| use burn::nn::PaddingConfig1d; | ||
| use burn::nn::pool::AvgPool1d; | ||
| use burn::nn::pool::AvgPool1dConfig; | ||
|
|
||
| #[derive(Module, Debug)] | ||
| pub struct Model <B: Backend> { | ||
| avg_pool1d: AvgPool1d, | ||
| phantom: core::marker::PhantomData<B>, | ||
| device: burn::module::Ignored<B::Device>, | ||
| } | ||
|
|
||
| impl<B: Backend> Model <B> { | ||
| #[allow(unused_variables)] | ||
| pub fn new(device: &B::Device) -> Self { | ||
| let avg_pool1d = AvgPool1dConfig::new(3) | ||
| .with_stride(1) | ||
| .with_padding(PaddingConfig1d::Valid) | ||
| .with_count_include_pad(true) | ||
| .init(); | ||
|
|
||
| Self { | ||
| avg_pool1d, | ||
| phantom: core::marker::PhantomData, | ||
| device: burn::module::Ignored(device.clone()), | ||
| } | ||
| } | ||
| #[allow(clippy::let_and_return, clippy::approx_constant)] | ||
| pub fn forward(&self, input: Tensor<B, 3>) -> Tensor<B, 3> { | ||
| let output = self.avg_pool1d.forward(input); | ||
|
|
||
| output | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| assert_tokens(graph.codegen(), expected); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.