Fix more minor issues with convolution construction and group norm.

liuliu · Jan 22, 2025 · c0bd3a5 · c0bd3a5
1 parent e27298e
commit c0bd3a5
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 7 deletions.
diff --git a/lib/nnc/ccv_cnnp_model_addons.c b/lib/nnc/ccv_cnnp_model_addons.c
@@ -1108,7 +1108,23 @@ static void _ccv_cnnp_convolution_build(ccv_cnnp_model_t* const super, ccv_nnc_s
 	if (self->format)
 		weights_params.format = self->format;
 	ccv_nnc_tensor_set_n(&weights_params, self->filters);
-	const int c = ccv_nnc_tensor_get_c(params);
+	const int a_nd = ccv_nnc_tensor_nd(params.dim);
+	int c;
+	switch (params.format)
+	{
+		case CCV_TENSOR_FORMAT_NHWC:
+			c = params.dim[a_nd - 1];
+			break;
+		case CCV_TENSOR_FORMAT_NCHW:
+			if (a_nd == k_nd + 1)
+				c = params.dim[0];
+			else
+				c = params.dim[a_nd <= 1 ? 0 : 1];
+			break;
+		case CCV_TENSOR_FORMAT_CHWN:
+			c = params.dim[0];
+			break;
+	}
 	assert(c % self->groups == 0);
 	ccv_nnc_tensor_set_c(&weights_params, nd, c / self->groups);
 	int hw = -1;

diff --git a/lib/nnc/cmd/norm/gpu/ccv_nnc_group_norm_gpu_cudnn.cu b/lib/nnc/cmd/norm/gpu/ccv_nnc_group_norm_gpu_cudnn.cu
@@ -88,12 +88,13 @@ static int _ccv_nnc_group_norm_forw(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_
 	assert(ccv_nnc_tensor_view_check_dim((ccv_nnc_tensor_view_t*)outputs[0], adim));
 	int x;
 	int n = 1;
-	for (x = 0; x < CCV_NNC_MAX_DIM + 2; x++)
+	const int a_nd = ccv_nnc_tensor_nd(adim);
+	for (x = 0; x < a_nd; x++)
 		n *= adim[x];
-	for (x = 0; x < CCV_NNC_MAX_DIM + 2; x++)
+	for (x = 0; x < a_nd; x++)
 		n /= rdim[x];
 	int rcount = 1;
-	for (x = 0; x < CCV_NNC_MAX_DIM + 2; x++)
+	for (x = 0; x < a_nd; x++)
 		rcount *= rdim[x];
 	const float inv_n = 1. / n;
 	cudnnReduceTensorDescriptor_t reduce = ccv_nnc_stream_context_get_reduce_tensor_descriptor(stream_context);
@@ -197,12 +198,13 @@ static int _ccv_nnc_group_norm_back(const ccv_nnc_cmd_t cmd, const ccv_nnc_hint_
 	static const float one = 1, zero = 0, neg_one = -1;
 	int x;
 	int n = 1;
-	for (x = 0; x < CCV_NNC_MAX_DIM + 2; x++)
+	const int g_nd = ccv_nnc_tensor_nd(gdim);
+	for (x = 0; x < g_nd; x++)
 		n *= gdim[x];
-	for (x = 0; x < CCV_NNC_MAX_DIM + 2; x++)
+	for (x = 0; x < g_nd; x++)
 		n /= rdim[x];
 	int gcount = 1, rcount = 1;
-	for (x = 0; x < CCV_NNC_MAX_DIM + 2; x++)
+	for (x = 0; x < g_nd; x++)
 		gcount *= gdim[x], rcount *= rdim[x];
 	const float neg_inv_n = -1. / n;
 	cudnnReduceTensorDescriptor_t reduce = ccv_nnc_stream_context_get_reduce_tensor_descriptor(stream_context);