126 lines
3.0 KiB
C
126 lines
3.0 KiB
C
#include <rtthread.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
#include <float.h>
|
|
#include <math.h>
|
|
|
|
#include "mnist.h"
|
|
#include "onnx.h"
|
|
|
|
int mnist(int argc, char const *argv[])
|
|
{
|
|
int img_index = 0;
|
|
if(argc == 2)
|
|
{
|
|
img_index = atoi(argv[1]);
|
|
}
|
|
print_img(img[img_index]);
|
|
|
|
// 1. Conv2D
|
|
int64_t shapeW3[] = {2, 1, 3, 3};
|
|
int64_t dimW3 = 4;
|
|
int64_t permW3_t[] = { 0, 2, 3, 1};
|
|
float* W3_t = transpose(W3, shapeW3, dimW3, permW3_t);
|
|
|
|
float* conv1 = (float*) malloc(sizeof(float)*28*28*2);
|
|
memset(conv1, 0, sizeof(sizeof(float)*28*28*2));
|
|
conv2D(img[img_index], 28, 28, 1, W3, 2, 3, 3, 1, 1, 1, 1, B3, conv1, 28, 28);
|
|
|
|
free(W3_t);
|
|
|
|
// 2. Relu
|
|
float* relu1 = (float*) malloc(sizeof(float)*28*28*2);
|
|
relu(conv1, 28*28*2, relu1);
|
|
|
|
free(conv1);
|
|
|
|
// 3. Maxpool
|
|
float* maxpool1 = (float*) malloc(sizeof(float)*14*14*2);
|
|
memset(maxpool1, 0, sizeof(sizeof(float)*14*14*2));
|
|
maxpool(relu1, 28, 28, 2, 2, 2, 0, 0, 2, 2, 14, 14, maxpool1);
|
|
|
|
free(relu1);
|
|
|
|
// 4. Conv2D
|
|
int64_t shapeW2[] = {2, 2, 3, 3};
|
|
int64_t dimW2 = 4;
|
|
int64_t perm_t[] = { 0, 2, 3, 1};
|
|
float* W2_t = transpose(W2, shapeW2, dimW2, perm_t);
|
|
|
|
float* conv2 = (float*) malloc(sizeof(float)*14*14*2);
|
|
memset(conv2, 0, sizeof(sizeof(float)*14*14*2));
|
|
conv2D(maxpool1, 14, 14, 2, W2_t, 2, 3, 3, 1, 1, 1, 1, B2, conv2, 14, 14);
|
|
|
|
free(W2_t);
|
|
free(maxpool1);
|
|
|
|
// 5. Relu
|
|
float* relu2 = (float*) malloc(sizeof(float)*14*14*2);
|
|
relu(conv2, 14*14*2, relu2);
|
|
|
|
free(conv2);
|
|
|
|
// 6. Maxpool
|
|
float* maxpool2 = (float*) malloc(sizeof(float)*7*7*2);
|
|
memset(maxpool2, 0, sizeof(sizeof(float)*7*7*2));
|
|
maxpool(relu2, 14, 14, 2, 2, 2, 0, 0, 2, 2, 7, 7, maxpool2);
|
|
|
|
free(relu2);
|
|
|
|
// Flatten NOT REQUIRED
|
|
|
|
// 7. Dense
|
|
int64_t shapeW1[] = {98, 4};
|
|
int64_t dimW1 = 2;
|
|
int64_t permW1_t[] = { 1, 0};
|
|
float* W1_t = transpose(W1, shapeW1, dimW1, permW1_t);
|
|
|
|
float* dense1 = (float*) malloc(sizeof(float)*4);
|
|
memset(dense1, 0, sizeof(sizeof(float)*4));
|
|
dense(maxpool2, W1_t, 98, 4, B1, dense1);
|
|
|
|
free(W1_t);
|
|
free(maxpool2);
|
|
|
|
// 8. Dense
|
|
int64_t shapeW[] = {4, 10};
|
|
int64_t dimW = 2;
|
|
int64_t permW_t[] = { 1, 0};
|
|
float* W_t = transpose(W, shapeW, dimW, permW_t);
|
|
|
|
float* dense2 = (float*) malloc(sizeof(float)*10);
|
|
memset(dense2, 0, sizeof(sizeof(float)*10));
|
|
dense(dense1, W_t, 4, 10, B, dense2);
|
|
|
|
free(W_t);
|
|
free(dense1);
|
|
|
|
// 9. Softmax
|
|
float* output = (float*) malloc(sizeof(float)*10);
|
|
memset(output, 0, sizeof(sizeof(float)*10));
|
|
softmax(dense2, 10, output);
|
|
|
|
// 10. Result
|
|
float max = 0;
|
|
int max_index = 0;
|
|
printf("\nPredictions: \n");
|
|
for(int i = 0; i < 10; i++)
|
|
{
|
|
printf("%f ", output[i]);
|
|
if(output[i] > max)
|
|
{
|
|
max = output[i];
|
|
max_index = i;
|
|
}
|
|
}
|
|
printf("\n");
|
|
printf("\nThe number is %d\n", max_index);
|
|
|
|
free(dense2);
|
|
free(output);
|
|
|
|
return 0;
|
|
}
|
|
MSH_CMD_EXPORT(mnist, mnist simple example)
|