saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

SizeInBytes Fix for Strided Tensor (#11224)

Browse source 
* SizeInBytes Fix for Strided Tensor

* resolve comments
This commit is contained in:
Vincent Wang 2022-04-19 15:13:00 +08:00 committed by GitHub
parent 3dac66698b
commit 06026fe8e6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 40 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
onnxruntime/core/framework/tensor.cc
View file

@ -60,8 +60,24 @@ void Tensor::InitOrtValue(MLDataType p_type, const TensorShape& shape, void* p_d
}
size_t Tensor::SizeInBytes() const {
#ifdef ENABLE_TRAINING
int64_t size = 1;
if (!IsContiguous()) {
for (size_t dim = 0; dim < shape_.NumDimensions(); ++dim) {
if (shape_[dim] == 0) {
size = 0;
break;
}
size += strides_[dim] * (shape_[dim] - 1);
}
} else {
size = shape_.Size();
}
#else
int64_t size = shape_.Size();
#endif
size_t ret;
if (!IAllocator::CalcMemSizeForArray(SafeInt<size_t>(shape_.Size()), dtype_->Size(), &ret)) {
if (!IAllocator::CalcMemSizeForArray(SafeInt<size_t>(size), dtype_->Size(), &ret)) {
ORT_THROW("tensor size overflow");
}
return ret;

17
onnxruntime/test/framework/tensor_test.cc
View file

@ -220,20 +220,33 @@ TEST(TensorTest, Strided) {
TensorShapeVector strides{12, 4, 1};
ASSERT_EQ(t.Shape(), shape);
ASSERT_THAT(t.Strides(), testing::ContainerEq(gsl::make_span(strides.cbegin(), strides.cend())));
ASSERT_EQ(t.SizeInBytes(), sizeof(float) * 24);
TensorShape new_shape({4, 2, 3});
TensorShapeVector new_strides{1, 12, 4};
t.SetShapeAndStrides(new_shape, new_strides);
EXPECT_FALSE(t.IsContiguous());
ASSERT_EQ(t.Shape(), new_shape);
ASSERT_THAT(t.Strides(), testing::ContainerEq(gsl::make_span(new_strides.cbegin(), new_strides.cend())));
ASSERT_EQ(t.SizeInBytes(), sizeof(float) * 24);
Tensor t2(DataTypeImpl::GetType<float>(), new_shape, data, alloc->Info(), 0L, gsl::make_span(new_strides));
EXPECT_FALSE(t2.IsContiguous());
ASSERT_EQ(t.Shape(), new_shape);
ASSERT_EQ(t2.Shape(), new_shape);
ASSERT_THAT(t2.Strides(), testing::ContainerEq(gsl::make_span(new_strides.cbegin(), new_strides.cend())));
ASSERT_EQ(t2.SizeInBytes(), sizeof(float) * 24);
t2.SetShapeAndStrides(shape, strides);
EXPECT_TRUE(t2.IsContiguous());
ASSERT_EQ(t.Shape(), new_shape);
ASSERT_EQ(t2.Shape(), shape);
ASSERT_THAT(t2.Strides(), testing::ContainerEq(gsl::make_span(strides.cbegin(), strides.cend())));
ASSERT_EQ(t2.SizeInBytes(), sizeof(float) * 24);
alloc->Free(data);
data = alloc->Alloc(sizeof(int64_t));
TensorShapeVector single_element_strides{0, 0, 0};
Tensor t3(DataTypeImpl::GetType<int64_t>(), shape, data, alloc->Info(), 0L, gsl::make_span(single_element_strides));
EXPECT_FALSE(t3.IsContiguous());
ASSERT_EQ(t3.Shape(), shape);
ASSERT_THAT(t3.Strides(),
testing::ContainerEq(gsl::make_span(single_element_strides.cbegin(), single_element_strides.cend())));
ASSERT_EQ(t3.SizeInBytes(), sizeof(int64_t));
alloc->Free(data);
}
#endif

8
orttraining/orttraining/eager/test/ort_ops.py
View file

@ -128,5 +128,13 @@ class OrtOpTests(unittest.TestCase):
ort_narrow = cpu_narrow.to('ort')
assert torch.allclose(cpu_narrow, ort_narrow.cpu())
def test_zero_stride(self):
print('ssssss')
device = self.get_device()
t = torch.empty_strided(size=(6, 1024, 512), stride=(0, 0, 0))
assert(t.storage().size() == 1) # This test is trying to confirm that transferring a tensor with a storage size of 1 works
ort_t = t.to(device)
assert torch.allclose(t, ort_t.cpu())
if __name__ == '__main__':
unittest.main()