1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
|
#pragma once
#include <ATen/core/functional.h>
#include <torch/csrc/WindowsTorchApiMacro.h>
#include <ATen/ATen.h>
#include <utility>
namespace torch { namespace utils {
inline at::Tensor flatten_dense_tensors(at::TensorList tensors) {
static auto flatten = [](const at::Tensor &t) { return t.contiguous().view({-1}); };
if (tensors.size() == 1)
return flatten(tensors[0]);
return at::cat(fmap(tensors, flatten));
}
inline std::vector<at::Tensor> unflatten_dense_tensors(const at::Tensor& flat, at::TensorList tensors) {
std::vector<at::Tensor> outputs;
outputs.reserve(tensors.size());
size_t offset = 0;
for (const auto & tensor : tensors) {
auto numel = tensor.numel();
outputs.push_back(flat.narrow(0, offset, numel).view(tensor.sizes()));
offset += numel;
}
return outputs;
}
struct TensorGroup {
std::vector<at::Tensor> tensors;
size_t size = 0;
at::DeprecatedTypeProperties& type() {
AT_ASSERT(!tensors.empty());
return tensors[0].type();
}
};
// Helper function that takes a list of tensors and splits them into tensor
// groups by the size limit and outputs these tensor groups. If the input
// tensors are of different tensor types, they will be split into different
// groups as well.
//
// Two options of splitting provided to the user,
//
// Imagine the size_limit is 256 and the list of input tensors are:
// tensor_a(fp16 - 128 bytes),
// tensor_b(fp32 - 256 bytes),
// tensor_c(fp16 - 128 bytes),
//
// when fine_grained == false:
// The function will read the list of tensors sequentially and accumulate
// enough tensors for each data type until the size_limit, therefore:
// it will output: {{tensor_a, tensor_c}, {tensor_b}}
//
// when fine_grained == true:
// The function will read the list of tensors sequentially and accumulate
// enough tensors for all data types until the size_limit, and then split
// the accumulated tensors into different groups by data types, therefore:
// it will output: {{tensor_a}, {tensor_b}, {tensor_c}}
TORCH_API std::vector<TensorGroup> take_tensors(
at::TensorList tensors,
size_t size_limit,
bool fine_grained = false);
TORCH_API void reorder_tensors_like(std::vector<at::Tensor>& tensors, at::TensorList order);
TORCH_API std::pair<at::Tensor, at::Tensor> flatten_sparse_tensors(at::TensorList tensors);
TORCH_API std::vector<at::Tensor> unflatten_sparse_tensors(
const at::Tensor& flat_indices,
const at::Tensor& flat_values,
at::TensorList tensors);
}}
|
