Xây dựng lớp mở rộng đạo hàm tùy biến kiểu PyTorch (Custom Autograd Function)
Trong PyTorch, khi muốn viết một toán tử mới tối ưu hóa (như FlashAttention, Gated Linear Unit, hay Quantized Activation), lập trình viên kế thừa lớp torch.autograd.Function với đố…
Tiến độ của tôi ở bài này
Điểm được lưu vào tài khoản sau khi chấm bài.
Đang tải điểm của bạn…
Kiến thức và chủ đề
Kiến thức tiên quyết: scalar-autograd-value-engine, vector-jacobian-product-vjp-autograd.
Nội dung đề bài
Mô tả bài toán
Trong PyTorch, khi muốn viết một toán tử mới tối ưu hóa (như FlashAttention, Gated Linear Unit, hay Quantized Activation), lập trình viên kế thừa lớp torch.autograd.Function với đối tượng lưu ngữ cảnh ctx:
ctx.save_for_backward(*tensors): lưu các tensor cần thiết trong forward để dùng lại trong backward.forward(ctx, *inputs): phương thức static nhận đầu vào, tính toán kết quả và lưu context.backward(ctx, *grad_outputs): phương thức static nhận gradient đầu ra, khôi phục tensor quactx.saved_tensors, và trả về gradient cho từng đầu vào tương ứng.
Hãy xây dựng framework AutogradFunction và cài đặt toán tử nhân ma trận MatMulFunction:
- Lớp
Context: - Phương thức
save_for_backward(self, *tensors): lưu các tensor vào thuộc tínhself.saved_tensors. - Lớp
MatMulFunction: forward(ctx: Context, A: np.ndarray, B: np.ndarray) -> np.ndarray:- Kiểm tra kích thước nhân ma trận hợp lệ (A ∈ RM × K, B ∈ RK × N).
- Lưu A, B vào
ctx.save_for_backward(A, B). - Trả về C = A B.
backward(ctx: Context, grad_output: np.ndarray) -> tuple[np.ndarray, np.ndarray]:- Nhận dC = grad_output ∈ RM × N.
- Lấy A, B = ctx.savedtensors.
- Tính dA = dC BT ∈ RM × K và dB = AT dC ∈ RK × N.
- Trả về tuple
(dA, dB).
Input
- Lớp
Context(): Khởi tạo đối tượng với các tham số, trọng số hoặc cấu hình tương ứng. - Lớp
MatMulFunction(): Khởi tạo đối tượng với các tham số, trọng số hoặc cấu hình tương ứng.
Output
- Các phương thức của
Context: Trả về kết quả tính toán hoặc cập nhật trạng thái nội bộ của đối tượng. - Các phương thức của
MatMulFunction: Trả về kết quả tính toán hoặc cập nhật trạng thái nội bộ của đối tượng.
Ràng buộc
- Thời gian chạy tối đa: 6000ms.
- Giới hạn bộ nhớ: 512MB.
- Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.
Ví dụ 1
Input
ctx = Context()
A = np.array([[1.0, 2.0], [3.0, 4.0]])
B = np.array([[5.0, 6.0], [7.0, 8.0]])
C = MatMulFunction.forward(ctx, A, B)
dC = np.ones((2, 2))
dA, dB = MatMulFunction.backward(ctx, dC)Output
([[11., 15.],
[11., 15.]], [[4., 4.],
[6., 6.]])Giải thích
Hàm/lớp được gọi với các tham số mẫu trên và trả về kết quả số học / kích thước tensor tương ứng theo đúng thiết kế.
Ví dụ 2
Input
ctx = Context()
A = np.random.randn(3, 4)
B = np.random.randn(4, 5)
C = MatMulFunction.forward(ctx, A, B)
dC = np.random.randn(3, 5)
dA, dB = MatMulFunction.backward(ctx, dC)Output
([[-0.1379, -1.0634, -1.2743, -2.3784],
[-0.0098, -1.5685, 1.3653, 0.0434],
[ 1.7467, 0.4037, 0.1793, -1.6243]], [[ 0.0401, 1.4466, -2.0129, -0.3987, -4.1517],
[-0.4688, -2.1962, 1.4287, 0.5913, 2.9233],
[ 0.5414, 1.4506, -0.2983, -0.2079, -0.611 ],
[-0.004 , 0.8536, -0.0135, -1.0419, 0.0851]])Giải thích
Hàm/lớp được gọi với các tham số mẫu trên và trả về kết quả số học / kích thước tensor tương ứng theo đúng thiết kế.
Gợi ý và lời giải chỉ mở sau khi bạn bấm Nộp bài. Giáo viên và quản trị viên mở được ngay.
Góp ý & báo lỗi bài tập
Đề bài chưa rõ, test có vấn đề hay bạn có ý tưởng giúp bài tốt hơn? Gửi cho đội ngũ AI Empire nhé — mỗi góp ý đều được đọc.