Skip to content

Speed up printing of large tensors. #6876

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

Merged
gchanan merged 5 commits into from
Apr 24, 2018
Merged

Speed up printing of large tensors. #6876

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
11ad445
Speed up printing of large tensors.
gchanan Apr 23, 2018
0bab332
Merge remote-tracking branch 'origin/master' into faster_printing
gchanan Apr 23, 2018
df71e20
Fix flake8.
gchanan Apr 23, 2018
7897c6f
Merge remote-tracking branch 'origin/master' into faster_printing
gchanan Apr 23, 2018
a029193
Add else case.
gchanan Apr 23, 2018
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
30 changes: 24 additions & 6 deletions torch/_tensor_str.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,8 +92,8 @@ def _number_format(tensor, min_sz=-1):

int_mode = True
# TODO: use fmod?
for value in tensor:
if value != math.ceil(value.item()):
for value in tensor.tolist():
if value != math.ceil(value):
int_mode = False
break

Expand Down Expand Up @@ -155,11 +155,11 @@ def _vector_str(self, indent, fmt, scale, sz, summarize):
char_per_line = element_length * elements_per_line

if summarize and self.size(0) > 2 * PRINT_OPTS.edgeitems:
data = ([fmt.format(val.item() / scale) for val in self[:PRINT_OPTS.edgeitems]] +
data = ([fmt.format(val / scale) for val in self[:PRINT_OPTS.edgeitems].tolist()] +
[' ...'] +
[fmt.format(val.item() / scale) for val in self[-PRINT_OPTS.edgeitems:]])
[fmt.format(val / scale) for val in self[-PRINT_OPTS.edgeitems:].tolist()])
else:
data = [fmt.format(val.item() / scale) for val in self]
data = [fmt.format(val) for val in self.tolist()]

data_lines = [data[i:i + elements_per_line] for i in range(0, len(data), elements_per_line)]
lines = [', '.join(line) for line in data_lines]
Expand Down Expand Up @@ -187,6 +187,24 @@ def _tensor_str(self, indent, fmt, scale, sz, summarize):
return '[' + tensor_str + ']'


def get_summarized_data(self):
dim = self.dim()
if dim == 0:
return self
if dim == 1:
if self.size(0) > 2 * PRINT_OPTS.edgeitems:
return torch.cat((self[:PRINT_OPTS.edgeitems], self[-PRINT_OPTS.edgeitems:]))
else:
return self
if self.size(0) > 2 * PRINT_OPTS.edgeitems:
start = [get_summarized_data(self[i]).view(-1) for i in range(0, PRINT_OPTS.edgeitems)]
end = ([get_summarized_data(self[i]).view(-1)
for i in range(len(self) - PRINT_OPTS.edgeitems, len(self))])
return torch.cat((start + end))
else:
return self


def _str(self):
if self.is_sparse:
size_str = str(tuple(self.shape)).replace(' ', '')
Expand Down Expand Up @@ -215,7 +233,7 @@ def _str(self):
if self.dtype != torch.get_default_dtype() and self.dtype != torch.int64:
suffix = ', dtype=' + str(self.dtype) + suffix

fmt, scale, sz = _number_format(self)
fmt, scale, sz = _number_format(get_summarized_data(self) if summarize else self)
if scale != 1:
prefix = prefix + SCALE_FORMAT.format(scale) + ' ' * indent
tensor_str = _tensor_str(self, indent, fmt, scale, sz, summarize)
Expand Down