Tính kích thước của optimizer
Bạn đang khảo sát các optimizer khác nhau để huấn luyện một mô hình và cần định lượng mức sử dụng bộ nhớ của optimizer để so sánh một cách khách quan. Làm thử, bạn đã tải một mô hình DistilBERT và optimizer AdamW để đo mức sử dụng bộ nhớ. Hãy viết hàm compute_optimizer_size để tính kích thước của một optimizer.
optimizer AdamW đã được định nghĩa trực tiếp (không dùng Trainer), và quá trình huấn luyện đã hoàn tất.
Bài tập này là một phần của khóa học
Huấn luyện Mô hình AI Hiệu quả với PyTorch
Hướng dẫn bài tập
- Tính số phần tử và kích thước của mỗi
tensortrong vòng lặpfor. - Tính tổng kích thước của
optimizertheo megabyte. - Truy cập state dictionary của optimizer bằng phương thức phù hợp trên
optimizer.state.
Bài tập tương tác thực hành trực tiếp
Hãy thử làm bài tập này bằng cách hoàn thành đoạn mã mẫu này.
def compute_optimizer_size(optimizer_state):
total_size_megabytes, total_num_elements = 0, 0
for params in optimizer_state:
for name, tensor in params.items():
tensor = torch.tensor(tensor)
# Compute number of elements and size of each tensor
num_elements, element_size = tensor.____(), tensor.____()
total_num_elements += num_elements
# Compute the total size in megabytes
total_size_megabytes += ____ * ____ / (1024 ** 2)
return total_size_megabytes, total_num_elements
# Pass in the optimizer state
total_size_megabytes, total_num_elements = compute_optimizer_size(optimizer.state.____())
print(f"Number of optimizer parameters: {total_num_elements:,}\nOptimizer size: {total_size_megabytes:.0f} MB")