ADD: ported first two layers of the fast-depth model

example/cnn/main.c

@@ -38,34 +38,63 @@ typedef struct {
   Tensor *conv1_1_running_mean;
   Tensor *conv1_1_running_var;
   Tensor *conv1_1_out;
+  Tensor *conv1_3_weight;
+  Tensor *conv1_3_out;
+  Tensor *conv1_4_weight;
+  Tensor *conv1_4_bias;
+  Tensor *conv1_4_running_mean;
+  Tensor *conv1_4_running_var;
+  Tensor *conv1_4_out;
 } Model;
 
 /**
  * Initialize the required tensors for the model
  */
 void init(Model *model) {
-  uint8_t *array_pointer = weights_data;
+  float *array_pointer = (float *)weights_data;
 
   model->input = NN_ones(4, (size_t[]){1, 3, 224, 224}, DTYPE_F32);
 
   model->conv0_0_weight = NN_tensor(4, (size_t[]){16, 3, 3, 3}, DTYPE_F32, array_pointer);
-  array_pointer += 16 * 3 * 3 * 3 * sizeof(float);
+  array_pointer += 16 * 3 * 3 * 3;
   model->conv0_0_out = NN_tensor(4, (size_t[]){1, 16, 112, 112}, DTYPE_F32, NULL);
 
   model->conv0_1_weight = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
-  // model->conv0_1_weight = NN_ones(1, (size_t[]){16, }, DTYPE_F32);
-  array_pointer += 16 * sizeof(float);
+  array_pointer += 16;
   model->conv0_1_bias = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
-  // model->conv0_1_bias = NN_zeros(1, (size_t[]){16, }, DTYPE_F32);
-  array_pointer += 16 * sizeof(float);
+  array_pointer += 16;
   model->conv0_1_running_mean = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
-  // model->conv0_1_running_mean = NN_ones(1, (size_t[]){16, }, DTYPE_F32);
-  array_pointer += 16 * sizeof(float);
+  array_pointer += 16;
   model->conv0_1_running_var = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
-  // model->conv0_1_running_var = NN_zeros(1, (size_t[]){16, }, DTYPE_F32);
-  array_pointer += 16 * sizeof(float);
+  array_pointer += 16;
   model->conv0_1_out = NN_tensor(4, (size_t[]){1, 16, 112, 112}, DTYPE_F32, NULL);
+
+  model->conv1_0_weight = NN_tensor(4, (size_t[]){16, 1, 3, 3}, DTYPE_F32, array_pointer);
+  array_pointer += 16 * 1 * 3 * 3;
+  model->conv1_0_out = NN_tensor(4, (size_t[]){1, 16, 112, 112}, DTYPE_F32, NULL);
+  model->conv1_1_weight = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
+  array_pointer += 16;
+  model->conv1_1_bias = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
+  array_pointer += 16;
+  model->conv1_1_running_mean = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
+  array_pointer += 16;
+  model->conv1_1_running_var = NN_tensor(1, (size_t[]){16, }, DTYPE_F32, array_pointer);
+  array_pointer += 16;
+  model->conv1_1_out = NN_tensor(4, (size_t[]){1, 16, 112, 112}, DTYPE_F32, NULL);
 
+  model->conv1_3_weight = NN_tensor(4, (size_t[]){56, 16, 1, 1}, DTYPE_F32, array_pointer);
+  array_pointer += 56 * 16 * 1 * 1;
+  model->conv1_3_out = NN_tensor(4, (size_t[]){1, 56, 112, 112}, DTYPE_F32, NULL);
+  model->conv1_4_weight = NN_tensor(1, (size_t[]){56, }, DTYPE_F32, array_pointer);
+  array_pointer += 56;
+  model->conv1_4_bias = NN_tensor(1, (size_t[]){56, }, DTYPE_F32, array_pointer);
+  array_pointer += 56;
+  model->conv1_4_running_mean = NN_tensor(1, (size_t[]){56, }, DTYPE_F32, array_pointer);
+  array_pointer += 56;
+  model->conv1_4_running_var = NN_tensor(1, (size_t[]){56, }, DTYPE_F32, array_pointer);
+  array_pointer += 56;
+  model->conv1_4_out = NN_tensor(4, (size_t[]){1, 56, 112, 112}, DTYPE_F32, NULL);
+
   size_t size = (size_t)weights_end - (size_t)weights_start;
   printf("weight size: %d\n", (int)size);
 
@@ -78,15 +107,35 @@ void init(Model *model) {
 void forward(Model *model) {
   NN_Conv2d_F32(
     model->conv0_0_out, model->input,
-    model->conv0_0_weight, NULL,
-    (size_t[]){3, 3}, (size_t[]){2, 2}, (size_t[]){1, 1}, 1
+    model->conv0_0_weight, NULL, (size_t[]){2, 2}, (size_t[]){1, 1}, 1
     );
   NN_BatchNorm2d_F32(
     model->conv0_1_out, model->conv0_0_out,
     model->conv0_1_weight, model->conv0_1_bias,
-    1e-5, 0.1, model->conv0_1_running_mean, model->conv0_1_running_var
+    1e-5, model->conv0_1_running_mean, model->conv0_1_running_var
     );
   NN_ReLU6Inplace_F32(model->conv0_1_out);
+
+  NN_Conv2d_F32(
+    model->conv1_0_out, model->conv0_1_out,
+    model->conv1_0_weight, NULL, (size_t[]){1, 1}, (size_t[]){1, 1}, 16
+    );
+  NN_BatchNorm2d_F32(
+    model->conv1_1_out, model->conv1_0_out,
+    model->conv1_1_weight, model->conv1_1_bias,
+    1e-5, model->conv1_1_running_mean, model->conv1_1_running_var
+    );
+  NN_ReLU6Inplace_F32(model->conv1_1_out);
+  NN_Conv2d_F32(
+    model->conv1_3_out, model->conv1_1_out,
+    model->conv1_3_weight, NULL, (size_t[]){1, 1}, (size_t[]){0, 0}, 1
+    );
+  NN_BatchNorm2d_F32(
+    model->conv1_4_out, model->conv1_3_out,
+    model->conv1_4_weight, model->conv1_4_bias,
+    1e-5, model->conv1_4_running_mean, model->conv1_4_running_var
+    );
+  NN_ReLU6Inplace_F32(model->conv1_4_out);
 }
 
 int main() {
@@ -102,7 +151,7 @@ int main() {
   // NN_printf(model->input);
 
   printf("output:\n");
-  NN_printf(model->conv0_1_out);
+  NN_printf(model->conv1_4_out);
 
   return 0;
 }

nn/inc/nn_batchnorm2d.h

@@ -15,14 +15,13 @@
  * @param weight: the learnable weights of the module of shape (channels), or NULL if no weight is applied
  * @param bias: the learnable bias of the module of shape (channels), or NULL if no bias is applied
  * @param eps: a value added to the denominator for numerical stability
- * @param momentum: the value used for the running_mean and running_var computation
  * @param running_mean: the running mean of the module of shape (channels), or NULL if no running mean is applied
  * @param running_var: the running variance of the module of shape (channels), or NULL if no running variance is applied
  */
 void NN_BatchNorm2d_F32(
     Tensor *out, Tensor *in,
     Tensor *weight, Tensor *bias,
-    float eps, float momentum, Tensor *running_mean, Tensor *running_va
+    float eps, Tensor *running_mean, Tensor *running_va
     );
 
 

nn/inc/nn_conv2d.h

@@ -16,15 +16,14 @@
  * @param bias: the learnable bias of the module of shape (channels_out), or NULL if no bias is applied
  * @param in_channels: number of channels in the input tensor
  * @param out_channels: number of channels produced by the convolution
- * @param kernel_size: size of the convolution kernel
  * @param stride: stride for the cross-correlation
  * @param padding: the amount of padding applied to the input
  * @param groups: number of blocked connections from input channels to output channels
  */
 void NN_Conv2d_F32(
   Tensor *out, Tensor *in, 
   Tensor *weight, Tensor *bias, 
-  size_t *kernel_size, size_t *stride, size_t *padding, size_t groups
+  size_t *stride, size_t *padding, size_t groups
   );
 
 

nn/src/batchnorm2d/nn_batchnorm2d.c

@@ -4,7 +4,7 @@
 void NN_BatchNorm2d_F32(
   Tensor *out, Tensor *in, 
   Tensor *weight, Tensor *bias,
-  float eps, float momentum, Tensor *running_mean, Tensor *running_var) {
+  float eps, Tensor *running_mean, Tensor *running_var) {
   assert(in->ndim == 4);
   assert(out->ndim == 4);
   assert(in->dtype == DTYPE_F32);

nn/src/conv2d/nn_conv2d.c

@@ -5,7 +5,7 @@
 void NN_Conv2d_F32(
   Tensor *out, Tensor *in, 
   Tensor *weight, Tensor *bias, 
-  size_t *kernel_size, size_t *stride, size_t *padding, size_t groups) {
+  size_t *stride, size_t *padding, size_t groups) {
   const size_t dilation[2] = {1, 1};
 
   assert(in->ndim == 4);
@@ -19,9 +19,9 @@ void NN_Conv2d_F32(
   }
   assert(out->shape[0] == in->shape[0]);
   assert(out->shape[1] == weight->shape[0]);
-  assert(in->shape[1] == weight->shape[1]);
-  assert(out->shape[2] == (in->shape[2] + 2 * padding[0] - dilation[0] * (kernel_size[0] - 1) - 1) / stride[0] + 1);
-  assert(out->shape[3] == (in->shape[3] + 2 * padding[1] - dilation[1] * (kernel_size[1] - 1) - 1) / stride[1] + 1);
+  assert(in->shape[1] == weight->shape[1] * groups);
+  assert(out->shape[2] == (in->shape[2] + 2 * padding[0] - dilation[0] * (weight->shape[2] - 1) - 1) / stride[0] + 1);
+  assert(out->shape[3] == (in->shape[3] + 2 * padding[1] - dilation[1] * (weight->shape[3] - 1) - 1) / stride[1] + 1);
   assert(groups > 0);
   assert(in->shape[1] % groups == 0);
   assert(out->shape[1] % groups == 0);
@@ -31,8 +31,8 @@ void NN_Conv2d_F32(
   size_t in_channels = in->shape[1];
   size_t input_height = in->shape[2];
   size_t input_width = in->shape[3];
-  size_t kernel_height = kernel_size[0];
-  size_t kernel_width = kernel_size[1];
+  size_t kernel_height = weight->shape[2];
+  size_t kernel_width = weight->shape[3];
   size_t stride_height = stride[0];
   size_t stride_width = stride[1];
   size_t padding_height = padding[0];
@@ -46,11 +46,11 @@ void NN_Conv2d_F32(
 
   for (size_t n = 0; n < batch_size; n += 1) {
     for (size_t g = 0; g < groups; g += 1) {
-      for (size_t oc = 0; oc < out_channels; oc += 1) {
+      for (size_t oc = 0; oc < out_channels / groups; oc += 1) {
         for (size_t oh = 0; oh < output_height; oh += 1) {
           for (size_t ow = 0; ow < output_width; ow += 1) {
             float sum = 0.0;
-            for (size_t ic = 0; ic < in_channels; ic += 1) {
+            for (size_t ic = 0; ic < in_channels / groups; ic += 1) {
               for (size_t kh = 0; kh < kernel_height; kh += 1) {
                 for (size_t kw = 0; kw < kernel_width; kw += 1) {
                   size_t ih = oh * stride_height + kh - padding_height;
@@ -72,7 +72,7 @@ void NN_Conv2d_F32(
               }
             }
             if (bias != NULL) {
-              sum += ((float *)bias->data)[oc];
+              sum += ((float *)bias->data)[g * (out_channels / groups) + oc];
             }
             size_t out_idx = n * out_channels * output_height * output_width
                                + (g * (out_channels / groups) + oc) * output_height * output_width