Merge pull request #222 from sammsaski/master

Add support for FlattenLayer3dInto2d
verivital · Jun 10, 2024 · 388483d · 388483d
2 parents 693b5dc + cbddbbb
commit 388483d
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/code/nnv/engine/nn/layers/FlattenLayer.m b/code/nnv/engine/nn/layers/FlattenLayer.m
@@ -136,15 +136,19 @@
                 else
                     error('Invalid input.');
                 end 
-            elseif strcmp(obj.Type, 'nnet.onnx.layer.FlattenLayer') || strcmp(obj.Type, 'nnet.onnx.layer.FlattenInto2dLayer') 
+            elseif strcmp(obj.Type, 'nnet.onnx.layer.FlattenLayer') || strcmp(obj.Type, 'nnet.onnx.layer.FlattenInto2dLayer') || strcmp(obj.Type, 'nnet.onnx.layer.Flatten3dInto2dLayer')
                 % C-style flatten
                 if length(n) == 2
                     image = permute(image,[2 1]);
                     flatten_im = reshape(image, [1 1 n(1)*n(2)]);
                 elseif length(n) == 3
                     image = permute(image,[2 1 3]);
                     flatten_im = reshape(image, [1 1 n(1)*n(2)*n(3)]);
+                elseif length(n) == 4
+                    image = permute(image, [3 2 1 4]);
+                    flatten_im = reshape(image, [1 1 n(1)*n(2)*n(3)*n(4)]);
                 else
+                    fprintf("Invalid input image with size %d", length(n));
                     error('Invalid input image');
                 end 
             else
@@ -383,7 +387,8 @@
             % date: 6/25/2021
 
             if ~isa(layer, 'nnet.keras.layer.FlattenCStyleLayer') && ~isa(layer, 'nnet.cnn.layer.FlattenLayer') ...
-                    && ~isa(layer, 'nnet.onnx.layer.FlattenLayer') && ~isa(layer, 'nnet.onnx.layer.FlattenInto2dLayer')
+                    && ~isa(layer, 'nnet.onnx.layer.FlattenLayer') && ~isa(layer, 'nnet.onnx.layer.FlattenInto2dLayer') ...
+                    && ~isa(layer, 'nnet.onnx.layer.Flatten3dInto2dLayer')
                 error('Input is not a flatten layer');
             end
 

diff --git a/code/nnv/engine/utils/matlab2nnv.m b/code/nnv/engine/utils/matlab2nnv.m
@@ -136,7 +136,7 @@
 
     % Flatten Layer
     elseif isa(L, 'nnet.keras.layer.FlattenCStyleLayer') || isa(L, 'nnet.cnn.layer.FlattenLayer') || isa(L, 'nnet.onnx.layer.FlattenLayer') ...
-            || isa(L, 'nnet.onnx.layer.FlattenInto2dLayer')
+            || isa(L, 'nnet.onnx.layer.FlattenInto2dLayer') || isa(L, 'nnet.onnx.layer.Flatten3dInto2dLayer')
         Li = FlattenLayer.parse(L);
 
     % Sigmoid Layer (also referred to as logsig)