[CNNTrain] Add ConGradX

Author: runwangdl

Deeploy/CommonExtensions/OptimizationPasses/TopologyOptimizationPasses/LoweringOptimizationPasses.py

@@ -227,7 +227,7 @@ def _NCHWtoNHWC_fun(graph: gs.Graph, match: Match, name: str, default_channels_f
         tensorIn = node.inputs[0]
         tensorOut = node.outputs[0]
 
-        if node.op in ["RequantizedConv", "Conv"]:
+        if node.op in ["RequantizedConv", "Conv", "ConvGradX"]:
             spatialDims = len(node.inputs[1].shape) - 2
         elif node.op in ["MaxPool", "AveragePool", "AveragePoolGrad"]:
             spatialDims = len(node.attrs["kernel_shape"])
@@ -242,8 +242,9 @@ def _NCHWtoNHWC_fun(graph: gs.Graph, match: Match, name: str, default_channels_f
         permuteOut = _transformLayoutPermutation(len(tensorOut.shape), spatialDims, channels_first)
         graph.nodes.append(_prependTranspose(tensorOut, node, permuteOut))
 
-        if node.op in ["Conv", "RequantizedConv"]:
+        if node.op in ["Conv", "RequantizedConv", "ConvGradX"]:
             # In the case of Conv: [weights, opt. bias], RequantizedConv: [weights, mul, add, opt. shift]
+            # ConvGradX: [weight] (no bias)
             for tensor in node.inputs[1:]:
                 _transformLayoutConst(tensor, spatialDims, default_channels_first)
 
@@ -279,6 +280,15 @@ def __init__(self, default_channels_first: bool = True):
         super().__init__(graph, partial(_NCHWtoNHWC_fun, default_channels_first = default_channels_first), name)
 
 
+@contextagnostic
+class NCHWtoNHWCConvGradXPass(ReplaceSequentialPatternPass):
+
+    def __init__(self, default_channels_first: bool = True):
+        graph = _singleNodePattern(op = "ConvGradX")
+        name = "_NCHW_TO_NHWC_CONVGRADX_PASS"
+        super().__init__(graph, partial(_NCHWtoNHWC_fun, default_channels_first = default_channels_first), name)
+
+
 @contextagnostic
 class NCHWtoNHWCConvPass(ReplaceSequentialPatternPass):
 
@@ -383,6 +393,7 @@ def __init__(self, default_channels_first: bool = True):
             NCHWtoNHWCMaxPoolPass(default_channels_first),
             NCHWtoNHWCAveragePoolPass(default_channels_first),
             NCHWtoNHWCAveragePoolGradPass(default_channels_first),
+            NCHWtoNHWCConvGradXPass(default_channels_first),
             NCHWtoNHWCDwConvPass(default_channels_first),
             NCHWtoNHWCConvPass(default_channels_first),
         ]
@@ -398,6 +409,7 @@ def __init__(self, default_channels_first: bool = True):
             NCHWtoNHWCMaxPoolPass(default_channels_first),
             NCHWtoNHWCAveragePoolPass(default_channels_first),
             NCHWtoNHWCAveragePoolGradPass(default_channels_first),
+            NCHWtoNHWCConvGradXPass(default_channels_first),
             PULPNCHWtoNHWCDwConvPass(default_channels_first),
             NCHWtoNHWCConvPass(default_channels_first),
         ]

Deeploy/Targets/PULPOpen/Bindings.py

@@ -251,6 +251,13 @@
         ForkTransformer) for float_type in FloatDataTypes
 ]
 
+PULPFloatConvTrans2DBindings = [
+    NodeBinding(
+        ConvChecker([PointerClass(float32_t), PointerClass(float32_t)],
+                    [PointerClass(float32_t)]), FloatConvTemplate.referenceConvTrans2DTemplate,
+        ForkTransformer)
+]
+
 PULPRQSMatrixVecBindings = [
     NodeBinding(
         PULPLinearChecker([PointerClass(type1),
@@ -464,3 +471,4 @@
     NodeBinding(DequantChecker([PointerClass(int32_t)], [PointerClass(float32_t)]), DequantTemplate.referenceTemplate,
                 ForkTransformer),
 ]
+

Deeploy/Targets/PULPOpen/Parsers.py

@@ -462,3 +462,8 @@ def parseNodeCtxt(self,
             return ctxt, False
 
         return newCtxt, True
+
+class PULPConvTrans2DParser(PULPFPConv2DParser):
+
+    def __init__(self, noBiasHoisting = True):
+            super().__init__(noBiasHoisting)

Deeploy/Targets/PULPOpen/Platform.py

@@ -34,7 +34,7 @@
 from Deeploy.Targets.PULPOpen.Layers import PULPRQSConvLayer, PULPRQSGEMMLayer
 from Deeploy.Targets.PULPOpen.Parsers import PULPConv1DParser, PULPConv2DParser, PULPDWConv1DParser, \
     PULPDWConv2DParser, PULPFPConv2DParser, PULPFPDWConv2DParser, PULPGEMMParser, PULPMatrixVecParser, \
-    PULPTallGEMMParser
+    PULPTallGEMMParser, PULPConvTrans2DParser
 from Deeploy.Targets.PULPOpen.Templates import AllocateTemplate, FreeTemplate
 from Deeploy.Targets.PULPOpen.Tiler import PULPAddTilingReadyBindings, PULPConcatTilingReadyBindings, \
     PULPConv2DTilingReadyBindings, PULPDWConv2DTilingReadyBindings, PULPFlattenTilingReadyBindings, \
@@ -48,7 +48,8 @@
     PULPRQSMatrixVecTilingReadyBindings, PULPRQSTallGEMMTilingReadyBindings, PULPRQSTilingReadyBindings, \
     PULPSGDTilingReadyBindings, PULPSliceTilingReadyBindings, PULPSoftmaxCrossEntropyGradTilingReadyBindings, \
     PULPSoftmaxCrossEntropyTilingReadyBindings, PULPSoftmaxGradTilingReadyBindings, PULPSoftmaxTilingReadyBindings, \
-    PULPTransposeTilingReadyBindings, PULPUniformRQSTilingReadyBindings, PULPAveragePool2DTilingReadyBindings, PULPAveragePoolGrad2DTilingReadyBindings
+    PULPTransposeTilingReadyBindings, PULPUniformRQSTilingReadyBindings, PULPAveragePool2DTilingReadyBindings, \
+    PULPAveragePoolGrad2DTilingReadyBindings, PULPConvTrans2DTilingReadyBindings
 from Deeploy.Targets.PULPOpen.TopologyOptimizationPasses.Passes import PULPAddRequantMergePass, \
     PULPConvRequantMergePass, PULPGEMMRequantMergePass, PULPMatMulRequantMergePass
 
@@ -77,6 +78,9 @@
 Conv1DMapper = NodeMapper(PULPConv1DParser(), [PULPConv1DBinding])
 DWConv1DMapper = NodeMapper(PULPDWConv1DParser(), [PULPDWConv1DBinding])
 FPConv2DMapper = NodeMapper(PULPFPConv2DParser(), PULPConv2DTilingReadyBindings)
+
+ConvGradXMapper = NodeMapper(PULPConvTrans2DParser(), PULPConvTrans2DTilingReadyBindings)
+
 Conv2DMapper = NodeMapper(PULPConv2DParser(), PULPRQSConv2DTilingReadyBindings)
 FPDWConv2DMapper = NodeMapper(PULPFPDWConv2DParser(), PULPDWConv2DTilingReadyBindings)
 DWConv2DMapper = NodeMapper(PULPDWConv2DParser(), PULPRQSDWConv2DTilingReadyBindings)
@@ -113,6 +117,7 @@
 AveragePoolGrad2DMapper = NodeMapper(AveragePool2DParser(), PULPAveragePoolGrad2DTilingReadyBindings)
 PULPMapping = {
     'Conv': ConvLayer([FPConv2DMapper, FPDWConv2DMapper]),
+    'ConvGradX': ConvLayer([ConvGradXMapper]),
     'RequantizedConv': PULPRQSConvLayer([Conv2DMapper, DWConv2DMapper, Conv1DMapper, DWConv1DMapper]),
     'RequantizedGemm': PULPRQSGEMMLayer([MatrixVecMapper, TallGEMMMapper, GEMMMapper]),
     'Gemm': GEMMLayer([FloatGEMMMapper, GEMMDequantMapper]),

Deeploy/Targets/PULPOpen/Templates/FloatConvTemplate.py

@@ -158,3 +158,24 @@ def hoistTransientBuffers(self, ctxt: NetworkContext,
     ref_${data_out}_${data_out} += ${ch_im_out} * ${dim_im_out_x} * ${dim_im_out_y};
 }
 """)
+
+
+referenceConvTrans2DTemplate = NodeTemplate("""
+// 2D FP ConvTranspose HWC (Name: ${nodeName}, Op: ${nodeOp})
+${data_in_type.typeName} ref_${data_out}_${data_in} = ${data_in};
+${data_out_type.typeName} ref_${data_out}_${data_out} = ${data_out};
+for (uint32_t n=0; n<${batch}; ++n) {
+    PULP_ConvTrans2d_fp${data_in_type.referencedType.typeWidth}_fp${weight_type.referencedType.typeWidth}_fp${data_out_type.referencedType.typeWidth}_HWC(
+        ref_${data_out}_${data_in},
+        ${dim_im_in_y}, ${dim_im_in_x}, ${ch_im_in},
+        ${weight}, ${ch_im_out},
+        ${dim_kernel_y}, ${dim_kernel_x},
+        ${stride_y}, ${stride_x},
+        ref_${data_out}_${data_out},
+        ${padding_y_top}, ${padding_y_bottom}, ${padding_x_left}, ${padding_x_right}
+    );
+
+    ref_${data_out}_${data_in} += ${ch_im_in} * ${dim_im_in_x} * ${dim_im_in_y};
+    ref_${data_out}_${data_out} += ${ch_im_out} * ${dim_im_out_x} * ${dim_im_out_y};
+}
+""")