Compute storage in bytes for parameters, gradients, saved activations, and optimizer state from their tensor shapes and supplied byte widths. Gradient tensors have the same shapes as the parameters. Count no optimizer-state tensors for SGD, one per parameter for AdaGrad, and two per parameter for Adam.
N=i=1∏ndi M=NwHere n is a tensor's number of dimensions, i indexes its dimensions, d is each dimension's size, N is its element count, w is the supplied bytes per element, and M is its storage in bytes. An empty shape denotes a scalar with one element; a zero-sized dimension gives zero elements. An empty shape collection contains no tensors. Exclude allocator overhead and any additional master-weight copies.
Return a Python dictionary with exactly the keys parameters, gradients, activations, optimizer_state, and total. Every value is an exact Python integer byte count, and total is the sum of the four categories.
Input: param_shapes = [[4,3],[3]], param_bytes_per_element = 2, grad_bytes_per_element = 4, activation_shapes = [[2,4]], activation_bytes_per_element = 2, optimizer = "adam", optimizer_bytes_per_element = 4
Output: {"parameters":30,"gradients":60,"activations":16,"optimizer_state":120,"total":226}
Explanation: The parameter tensors contain 15 elements. Parameters use 30 bytes, gradients use 60, activations use 16, and Adam's two states use 120.
Input: param_shapes = [[2,2]], param_bytes_per_element = 4, grad_bytes_per_element = 4, activation_shapes = [], activation_bytes_per_element = 4, optimizer = "sgd", optimizer_bytes_per_element = 4
Output: {"parameters":16,"gradients":16,"activations":0,"optimizer_state":0,"total":32}
Input: param_shapes = [[128,64],[64]], param_bytes_per_element = 2, grad_bytes_per_element = 2, activation_shapes = [[32,128]], activation_bytes_per_element = 2, optimizer = "adagrad", optimizer_bytes_per_element = 4
Output: {"parameters":16512,"gradients":16512,"activations":8192,"optimizer_state":33024,"total":74240}
math.prod returns one for an empty shape.
Python integers avoid fixed-width overflow in large storage counts.
Accepts: array
Accepts: number
Accepts: number
Accepts: array
Accepts: number
Accepts: string
Accepts: number
Compute storage in bytes for parameters, gradients, saved activations, and optimizer state from their tensor shapes and supplied byte widths. Gradient tensors have the same shapes as the parameters. Count no optimizer-state tensors for SGD, one per parameter for AdaGrad, and two per parameter for Adam.
N=i=1∏ndi M=NwHere n is a tensor's number of dimensions, i indexes its dimensions, d is each dimension's size, N is its element count, w is the supplied bytes per element, and M is its storage in bytes. An empty shape denotes a scalar with one element; a zero-sized dimension gives zero elements. An empty shape collection contains no tensors. Exclude allocator overhead and any additional master-weight copies.
Return a Python dictionary with exactly the keys parameters, gradients, activations, optimizer_state, and total. Every value is an exact Python integer byte count, and total is the sum of the four categories.
Input: param_shapes = [[4,3],[3]], param_bytes_per_element = 2, grad_bytes_per_element = 4, activation_shapes = [[2,4]], activation_bytes_per_element = 2, optimizer = "adam", optimizer_bytes_per_element = 4
Output: {"parameters":30,"gradients":60,"activations":16,"optimizer_state":120,"total":226}
Explanation: The parameter tensors contain 15 elements. Parameters use 30 bytes, gradients use 60, activations use 16, and Adam's two states use 120.
Input: param_shapes = [[2,2]], param_bytes_per_element = 4, grad_bytes_per_element = 4, activation_shapes = [], activation_bytes_per_element = 4, optimizer = "sgd", optimizer_bytes_per_element = 4
Output: {"parameters":16,"gradients":16,"activations":0,"optimizer_state":0,"total":32}
Input: param_shapes = [[128,64],[64]], param_bytes_per_element = 2, grad_bytes_per_element = 2, activation_shapes = [[32,128]], activation_bytes_per_element = 2, optimizer = "adagrad", optimizer_bytes_per_element = 4
Output: {"parameters":16512,"gradients":16512,"activations":8192,"optimizer_state":33024,"total":74240}
math.prod returns one for an empty shape.
Python integers avoid fixed-width overflow in large storage counts.
Accepts: array
Accepts: number
Accepts: number
Accepts: array
Accepts: number
Accepts: string
Accepts: number
