Атрибут nbytes
Атрибут nbytes возвращает общее количество байтов, занимаемое тензором в памяти. Это значение вычисляется как произведение общего количества элементов тензора и размера одного элемента в байтах. Атрибут доступен только для чтения и полезен для оценки потребления памяти моделью или отдельными тензорами.
Синтаксис
tensor.nbytes
Пример
Давайте создадим целочисленный тензор и посмотрим, сколько байтов он занимает:
import torch
t = torch.tensor([1, 2, 3, 4, 5])
res = t.nbytes
print(res)
Результат выполнения кода:
20
По умолчанию целочисленный тензор использует тип int64, который занимает 8 байтов на элемент. В тензоре 5 элементов, поэтому общий размер составляет 40 байтов.
Пример
Давайте сравним размер тензора с разными типами данных:
import torch
t1 = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
t2 = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float64)
t3 = torch.tensor([1, 2, 3, 4, 5], dtype=torch.int16)
print(t1.nbytes)
print(t2.nbytes)
print(t3.nbytes)
Результат выполнения кода:
20
40
10
Тензор с типом float32 занимает 4 байта на элемент, float64 - 8 байтов, а int16 - 2 байта.
Пример
Давайте посмотрим, как размер тензора зависит от его формы:
import torch
t1 = torch.zeros(10, 10)
t2 = torch.zeros(100, 100)
t3 = torch.zeros(1000, 1000)
print(t1.nbytes)
print(t2.nbytes)
print(t3.nbytes)
Результат выполнения кода:
400
40000
4000000
Размер тензора растёт пропорционально количеству элементов. Тензор размером 1000x1000 с типом float32 занимает 4 мегабайта.
Пример
Давайте используем nbytes для оценки памяти модели:
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(784, 256),
nn.ReLU(),
nn.Linear(256, 10)
)
total_bytes = 0
for param in model.parameters():
total_bytes += param.nbytes
print(f"Model size: {total_bytes} bytes")
print(f"Model size: {total_bytes / 1024 / 1024:.2f} MB")
Результат выполнения кода:
"Model size: 803850 bytes"
"Model size: 0.77 MB"
Атрибут nbytes позволяет легко подсчитать общий размер всех параметров модели.