61 lines
1.1 KiB
C++
61 lines
1.1 KiB
C++
#pragma once
|
|
|
|
#include <cstdlib>
|
|
#include <vector>
|
|
|
|
// You can modify the data structure as you want
|
|
struct Tensor {
|
|
|
|
// Alloc memory
|
|
Tensor(std::vector<int> shape_) {
|
|
ndim = shape_.size();
|
|
for (int i = 0; i < ndim; i++) {
|
|
shape[i] = shape_[i];
|
|
}
|
|
int n = num_elem();
|
|
buf = (float*)malloc(n * sizeof(float));
|
|
}
|
|
|
|
// Alloc memory and copy
|
|
Tensor(std::vector<int> shape_, float *buf_) {
|
|
ndim = shape_.size();
|
|
for (int i = 0; i < ndim; i++) {
|
|
shape[i] = shape_[i];
|
|
}
|
|
int n = num_elem();
|
|
buf = (float*)malloc(n * sizeof(float));
|
|
for (int i=0; i<n; ++i){
|
|
buf[i] = buf_[i];
|
|
}
|
|
}
|
|
|
|
~Tensor() {
|
|
if (buf != nullptr)
|
|
free(buf);
|
|
}
|
|
|
|
void set_zero() {
|
|
int n = num_elem();
|
|
buf = (float*)malloc(n * sizeof(float));
|
|
for (int i = 0; i < n; i++){
|
|
buf[i] = 0.0;
|
|
}
|
|
}
|
|
|
|
int num_elem() {
|
|
int sz = 1;
|
|
for (int i = 0; i < ndim; i++)
|
|
sz *= shape[i];
|
|
return sz;
|
|
}
|
|
|
|
// Pointer to data
|
|
float *buf = nullptr;
|
|
|
|
// Shape of tensor, from outermost dimension to innermost dimension.
|
|
// e.g., {{1.0, -0.5, 2.3}, {4.3, 5.6, -7.8}} => shape = {2, 3}
|
|
int ndim = 0;
|
|
int shape[4];
|
|
};
|
|
|