Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TensorV2] Add functions to split and concatenate tensors @open sesame 03/07 10:42 #2497

Merged
merged 1 commit into from
Mar 11, 2024
Merged

[TensorV2] Add functions to split and concatenate tensors @open sesame 03/07 10:42 #2497

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
198 changes: 195 additions & 3 deletions nntrainer/tensor/float_tensor.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -162,12 +162,19 @@ const void *FloatTensor::getAddress(unsigned int i) const {
return &((float *)getData())[i];
}

const float FloatTensor::getValue(unsigned int i) const {
const float &FloatTensor::getValue(unsigned int i) const {
return ((float *)getData())[i];
}

const float FloatTensor::getValue(unsigned int b, unsigned int c,
unsigned int h, unsigned int w) const {
float &FloatTensor::getValue(unsigned int i) { return ((float *)getData())[i]; }

const float &FloatTensor::getValue(unsigned int b, unsigned int c,
unsigned int h, unsigned int w) const {
return getValue(getIndex(b, c, h, w));
}

float &FloatTensor::getValue(unsigned int b, unsigned int c, unsigned int h,
unsigned int w) {
return getValue(getIndex(b, c, h, w));
}

Expand Down Expand Up @@ -896,6 +903,191 @@ void FloatTensor::zoneout_mask(TensorV2 &opposite, float zoneout) {
}
}

std::vector<TensorV2> FloatTensor::split(std::vector<size_t> sizes, int axis) {
size_t num_size = sizes.size();

if (axis == -1) {
axis = 3;
}

size_t total_size = std::accumulate(sizes.begin(), sizes.end(), 0);
NNTR_THROW_IF(dim.getTensorDim(axis) != total_size, std::invalid_argument)
<< "given sum of sizes did not match with origin tensor dim, tensor dim: "
<< dim.getTensorDim(axis) << " total size: " << total_size;

std::vector<TensorDim> ret_dims;
ret_dims.reserve(num_size);
for (unsigned int i = 0; i < num_size; ++i) {
ret_dims[i] = dim;
ret_dims[i].setTensorDim(axis, sizes[i]);
}

bool is_format_nchw = (dim.getFormat() == Tformat::NCHW) ? true : false;
std::vector<TensorV2> ret;

auto iter_value = [this, is_format_nchw](
std::array<size_t, 4> &loc,
const std::array<size_t, 4> &end_loc,
const std::array<size_t, 4> &reset_dim_arr) -> float & {
auto &value = (is_format_nchw) ? getValue(loc[0], loc[1], loc[2], loc[3])
: getValue(loc[0], loc[3], loc[1], loc[2]);
for (int i = 3; i >= 0; --i) {
loc[i]++;
if (loc[i] == end_loc[i]) {
loc[i] -= reset_dim_arr[i];
continue;
}
break;
}
return value;
};

ret.reserve(num_size);

unsigned int accumulated_size = 0;
for (unsigned int i = 0; i < num_size; ++i) {
std::array<size_t, 4> loc = {0, 0, 0, 0};

if (is_format_nchw) {
loc[axis] += accumulated_size;
} else {
if (axis == 0) {
loc[0] += accumulated_size;
} else if (axis == 1) {
loc[3] += accumulated_size;
} else if (axis == 2 || axis == 3) {
loc[axis - 1] += accumulated_size;
}
}

ret.emplace_back(ret_dims[i]);
auto &ret_t = ret.back();

std::array<size_t, 4> end_loc;

if (is_format_nchw) {
end_loc = {ret_dims[i].batch(), ret_dims[i].channel(),
ret_dims[i].height(), ret_dims[i].width()};
} else {
end_loc = {ret_dims[i].batch(), ret_dims[i].height(), ret_dims[i].width(),
ret_dims[i].channel()};
}

accumulated_size += sizes[i];

if (is_format_nchw) {
end_loc[axis] = accumulated_size;
} else {
if (axis == 0) {
end_loc[0] = accumulated_size;
} else if (axis == 1) {
end_loc[3] = accumulated_size;
} else if (axis == 2 || axis == 3) {
end_loc[axis - 1] = accumulated_size;
}
}

std::array<size_t, 4> reset_dim_arr;
if (is_format_nchw) {
reset_dim_arr = {ret_dims[i].batch(), ret_dims[i].channel(),
ret_dims[i].height(), ret_dims[i].width()};
} else {
reset_dim_arr = {ret_dims[i].batch(), ret_dims[i].height(),
ret_dims[i].width(), ret_dims[i].channel()};
}

ret_t.apply_i<float>(
[&iter_value, &loc, &end_loc, &reset_dim_arr](float _) {
return iter_value(loc, end_loc, reset_dim_arr);
});
}

return ret;
}

TensorV2 FloatTensor::cat(const std::vector<TensorV2> &tensors, int axis) {
if (axis == -1) {
axis = 3;
}

TensorV2 ret;
auto ref_dim = tensors.front().getDim();
bool is_format_nchw = (ref_dim.getFormat() == Tformat::NCHW);
ref_dim.setTensorDim(axis, 1);
NNTR_THROW_IF(!std::all_of(tensors.begin(), tensors.end(),
[&ref_dim, axis](const TensorV2 &t) {
auto cur_dim = t.getDim();
cur_dim.setTensorDim(axis, 1);
return ref_dim == cur_dim;
}),
std::invalid_argument)
<< " all tensor must have the same dimension except for the axis, ref_dim: "
<< ref_dim << " axis : " << axis;

auto axis_dim = std::accumulate(tensors.begin(), tensors.end(), 0u,
[axis](unsigned cur, const TensorV2 &t) {
return cur += t.getDim().getTensorDim(axis);
});
auto iter_value =
[is_format_nchw](std::array<unsigned, 4> &loc,
const std::array<unsigned, 4> &start_loc, TensorV2 &t,
const std::array<unsigned, 4> &ref_dim_arr) -> float & {
auto &value = is_format_nchw
? t.getValue<float>(loc[0], loc[1], loc[2], loc[3])
: t.getValue<float>(loc[0], loc[3], loc[1], loc[2]);

for (int i = 3; i >= 0; --i) {
loc[i]++;
if (loc[i] - start_loc[i] == ref_dim_arr[i]) {
loc[i] = start_loc[i];
continue;
}
break;
}
return value;
};

auto ret_dim = ref_dim;
ret_dim.setTensorDim(axis, axis_dim);

ret = TensorV2(ret_dim);

std::array<unsigned, 4> loc = {0, 0, 0, 0};
for (auto &t : tensors) {
std::array<unsigned, 4> start_loc = loc;
std::array<unsigned, 4> tensor_dim_arr;
if (is_format_nchw) {
tensor_dim_arr[0] = t.getDim().getTensorDim(0);
tensor_dim_arr[1] = t.getDim().getTensorDim(1);
tensor_dim_arr[2] = t.getDim().getTensorDim(2);
tensor_dim_arr[3] = t.getDim().getTensorDim(3);
} else {
tensor_dim_arr[0] = t.getDim().getTensorDim(0);
tensor_dim_arr[1] = t.getDim().getTensorDim(2);
tensor_dim_arr[2] = t.getDim().getTensorDim(3);
tensor_dim_arr[3] = t.getDim().getTensorDim(1);
}

for (size_t i = 0u, sz = t.size(); i < sz; ++i) {
iter_value(loc, start_loc, ret, tensor_dim_arr) = t.getValue<float>(i);
}

if (is_format_nchw) {
loc[axis] += t.getDim().getTensorDim(axis);
} else {
if (axis == 0) {
loc[0] += t.getDim().getTensorDim(axis);
} else if (axis == 1) {
loc[3] += t.getDim().getTensorDim(axis);
} else if (axis == 2 || axis == 3) {
loc[axis - 1] += t.getDim().getTensorDim(axis);
}
}
}

return ret;
}

void FloatTensor::print(std::ostream &out) const {
printInstance(out, this);
const float *data = (float *)getData();
Expand Down
32 changes: 29 additions & 3 deletions nntrainer/tensor/float_tensor.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,13 @@ class FloatTensor : public TensorBase {
* @brief return value at specific location
* @param[in] i index
*/
const float getValue(unsigned int i) const;
const float &getValue(unsigned int i) const;

/**
* @brief return value at specific location
* @param[in] i index
*/
float &getValue(unsigned int i);

/**
* @brief return value at specific location
Expand All @@ -128,8 +134,18 @@ class FloatTensor : public TensorBase {
* @param[in] h height location
* @param[in] w width location
*/
const float getValue(unsigned int b, unsigned int c, unsigned int h,
unsigned int w) const;
const float &getValue(unsigned int b, unsigned int c, unsigned int h,
unsigned int w) const;

/**
* @brief return value at specific location
* @param[in] b batch location
* @param[in] c channel location
* @param[in] h height location
* @param[in] w width location
*/
float &getValue(unsigned int b, unsigned int c, unsigned int h,
unsigned int w);

/**
* @copydoc TensorV2::setValue(float value)
Expand Down Expand Up @@ -302,6 +318,16 @@ class FloatTensor : public TensorBase {
*/
void zoneout_mask(TensorV2 &opposite, float zoneout) override;

/**
* @copydoc TensorV2::split(std::vector<size_t> sizes, int axis)
*/
std::vector<TensorV2> split(std::vector<size_t> sizes, int axis) override;

/**
* @copydoc TensorV2::cat(const std::vector<TensorV2> &tensors, int axis)
*/
static TensorV2 cat(const std::vector<TensorV2> &tensors, int axis);

/**
* @copydoc TensorV2::copy(const TensorV2 &from)
*/
Expand Down
Loading
Loading