Fix Transpose Convolutions

deephealthproject · Jul 15, 2024 · fa99e51 · fa99e51
1 parent cc92a6b
commit fa99e51
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/descriptors/descriptor_convT2D.cpp b/src/descriptors/descriptor_convT2D.cpp
@@ -79,7 +79,8 @@ void ConvolDescriptorT2D::build(Tensor *A) {
 
     I = A;
 
-    nk = A->shape[1]; //ksize[0];
+    //nk = A->shape[1]; 
+    nk = ksize[0];
     kr = ksize[1];
     kc = ksize[2];
     kz = A->shape[1]/groups;
@@ -213,7 +214,7 @@ void ConvolDescriptorT2D::build(Tensor *A) {
                  in,iz,ir,ic);
    cudnnCreateFilterDescriptor(&wDesc);
    //CONVT we need to swap input channels with output so all other swappings (forward and backward functions) matches
-   cudnnSetFilter4dDescriptor(wDesc, data_type, tensor_format, nk, kz, kr, kc);
+   cudnnSetFilter4dDescriptor(wDesc, data_type, tensor_format, kz, nz, kr, kc);
 
    cudnnCreateTensorDescriptor(&yDesc);
    cudnnSetTensor4dDescriptor(yDesc, tensor_format, data_type, in, z,r,c);

diff --git a/src/serialization/onnx/net/layers/conv/convT_onnx.cpp b/src/serialization/onnx/net/layers/conv/convT_onnx.cpp
@@ -87,7 +87,8 @@ Layer* build_convT_layer(onnx::NodeProto *node,
   string weights_name = node->input(1); // Get weights and dims
   vector<float> *weights = &(map_init_values[weights_name]);
   vector<int> dims = map_init_dims[weights_name];
-  filters = dims[0];
+  //filters = dims[0];
+  filters = dims[1];
 
   // Deduce conv dimension from layer input
   if (parent_shape.size() == 3)