MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / determine_format

Method determine_format

tools/Polygraphy/polygraphy/util/format.py:68–98  ·  view source on GitHub ↗

Guesses the data format of a given shape. Args: shape (Tuple[int]): The shape, including batch dimension. Returns: DataFormat: The determined data format.

(shape)

Source from the content-addressed store, hash-verified

66
67 @staticmethod
68 def determine_format(shape):
69 """
70 Guesses the data format of a given shape.
71
72 Args:
73 shape (Tuple[int]): The shape, including batch dimension.
74
75 Returns:
76 DataFormat: The determined data format.
77 """
78 # The smaller this ratio, the closer a and b are.
79 def minmax_ratio(a, b):
80 return abs(max(a, b) / min(a, b))
81
82 # Assume all shapes include batch dimension
83 if len(shape) == 4:
84 # Typically, H and W are quite close, so if minmax_ratio(0, 1) > minmax_ratio(1, 2), then we assume CHW.
85 if minmax_ratio(shape[1], shape[2]) > minmax_ratio(shape[2], shape[3]):
86 return DataFormat.NCHW
87 return DataFormat.NHWC
88 elif len(shape) == 3:
89 return DataFormat.NHW
90 elif len(shape) == 2:
91 return DataFormat.NW
92 else:
93 G_LOGGER.warning(
94 "Cannot determine format for "
95 + str(shape)
96 + ". Currently only implemented for input_buffers with 1-3 non-batch dimensions. Please update this function!"
97 )
98 return DataFormat.UNKNOWN
99
100 # Get the permutation required to transpose old_format to new_format
101 @staticmethod

Callers 3

call_implMethod · 0.80
convertMethod · 0.80
test_format_deductionFunction · 0.80

Calls 1

warningMethod · 0.45

Tested by 1

test_format_deductionFunction · 0.64