Giải phóng khả năng tăng tốc gấp 2 lần cho DNN: Chuyển đổi mô hình với Mạng thần kinh tích hợp chỉ trong 1 phút

Mạng thần kinh tích hợp (INN) là các kiến trúc linh hoạt được đào tạo một lần và có thể chuyển đổi thành kích thước tùy ý do người dùng xác định mà không cần tinh chỉnh. Vì sóng âm thanh (âm nhạc) có thể được lấy mẫu ở bất kỳ tốc độ lấy mẫu (chất lượng âm thanh) mong muốn nào, INN có thể thay đổi hình dạng dữ liệu và thông số (chất lượng DNN) một cách linh hoạt.

Nhóm TheStage.ai đã trình bày bài báo của họ, “Mạng thần kinh tích hợp” tại hội nghị IEEE/CVF CVPR 2023. Tác phẩm này được công nhận là một trong 12 tác phẩm 'Ứng cử viên Giải thưởng' duy nhất được giới thiệu tại sự kiện. INN là một lớp mạng thần kinh mới kết hợp các tham số liên tục và toán tử tích phân để biểu diễn các lớp cơ bản. Trong giai đoạn suy luận, INN chuyển đổi thành biểu diễn DNN thông thường bằng cách lấy mẫu rời rạc từ các trọng số liên tục. Các tham số của các mạng như vậy là liên tục dọc theo các kích thước bộ lọc và kênh, dẫn đến việc cắt bớt có cấu trúc mà không tinh chỉnh chỉ bằng cách rời rạc hóa lại dọc theo các kích thước đó.

Trong bài viết này, chúng tôi sẽ hướng dẫn bạn quy trình chuyển đổi mô hình EDSR độ phân giải siêu hình ảnh 4x sang INN, sau đó trình bày cách đạt được việc cắt tỉa có cấu trúc của mô hình. Cuối cùng, chúng tôi sẽ chuyển đổi INN trở lại DNN rời rạc và triển khai nó trên GPU Nvidia để suy luận hiệu quả. Bài viết này sẽ tiến hành như sau:

1. Giới thiệu ngắn gọn về INN.
2. Tổng quan về mạng EDSR cho nhiệm vụ siêu phân giải.
3. Ứng dụng khung TorchIntegral để có được EDSR tích hợp trong một dòng mã.
4. Cắt tỉa cấu trúc của INN mà không tinh chỉnh INN (đường dẫn nhanh).
5. Triển khai các mô hình được cắt bớt trên GPU Nvidia.

Để biết thêm thông tin và cập nhật, vui lòng kiểm tra các tài nguyên sau:

Trang web dự án INN

Dự án INN Github

Mã hỗ trợ của bài viết này

INN để cắt tỉa DNN mà không cần tinh chỉnh

Các lớp trong INN được thay thế bằng toán tử tích phân, nhưng việc đánh giá thực tế các toán tử tích phân đòi hỏi sự rời rạc hóa tín hiệu đầu vào để sử dụng các phương pháp tích phân số. Hóa ra các lớp trong INN được thiết kế sao cho trùng khớp với các lớp DNN cổ điển (được kết nối đầy đủ, tích chập) sau khi rời rạc hóa.

Cắt tỉa hình ảnh siêu phân giải EDSR 4x

Siêu phân giải hình ảnh là một nhiệm vụ thị giác máy tính nổi tiếng, trong đó hình ảnh phải được nâng cao bằng toán tử suy giảm đã biết hoặc chưa biết. Chúng tôi xem xét hình thức siêu phân giải cổ điển bằng cách sử dụng phương pháp lấy mẫu hai khối làm toán tử suy giảm.

Nhiều kiến trúc có thể được sử dụng cho các tác vụ có độ phân giải siêu cao của hình ảnh, bao gồm các mạng thần kinh cao cấp dựa trên các mô hình khuếch tán và máy biến áp. Trong tài liệu này, chúng tôi sẽ tập trung vào kiến trúc EDSR 4x. Kiến trúc EDSR rất phù hợp cho phần trình diễn của chúng tôi vì nó bao gồm ResNet (được sử dụng rộng rãi trong nhiều vấn đề về học sâu) và khối lấy mẫu tăng 4x ở cuối. Một mô tả sơ đồ của EDSR có thể được tìm thấy trong các hình sau.

Đặc điểm của việc cắt tỉa kiến trúc EDSR

Cắt tỉa có cấu trúc, bao gồm việc xóa toàn bộ bộ lọc hoặc kênh, có ý nghĩa đặc biệt đối với các khối dư, đóng vai trò là khối xây dựng chính trong EDSR. Vì mỗi trạng thái được cập nhật bằng cách thêm khối Conv -> ReLU -> Conv vào đầu vào nên tín hiệu đầu vào và đầu ra phải có cùng số lượng tham số. Điều này có thể được quản lý hiệu quả trong TorchIntegral bằng cách tạo biểu đồ phụ thuộc cắt tỉa . Hình dưới đây minh họa rằng tích chập thứ hai của mỗi khối dư tạo thành một nhóm duy nhất.

Để tỉa bớt tích chập thứ hai trong khối dư, cần phải tỉa tất cả các tích chập thứ hai trong mỗi khối dư. Tuy nhiên, để thiết lập linh hoạt hơn, chúng tôi lược bỏ các bộ lọc của các tích chập đầu tiên và do đó các kênh của các tích chập thứ hai, trong tất cả các khối còn lại.

Chuyển đổi mô hình EDSR sang INN EDSR

Đối với các chuyển đổi từ DNN được đào tạo trước, chúng tôi sử dụng thuật toán hoán vị kênh lọc đặc biệt của mình với khả năng nội suy mượt mà hơn. Thuật toán hoán vị duy trì chất lượng mô hình trong khi làm cho trọng số của DNN trông giống như được lấy mẫu từ các hàm liên tục.

 import torch import torchintegral as inn from super_image import EdsrModel # creating 4x EDSR model model = EdsrModel.from_pretrained("eugenesiow/edsr", scale=4).cuda() # Transform model layers to integral. # continous_dims and discrete dims define which dimensions # of parameters tensors should be parametrized continuously # or stay fixed size as in discrete networks. # In our case we make all filter and channel dimensions # to be continuous excluding convolutions of the upsample block. model = inn.IntegralWrapper(init_from_discrete=True)( model, example_input, continuous_dims, discrete_dims ).cuda()

Điều chỉnh lưới tích hợp: cắt tỉa DNN sau đào tạo có cấu trúc

Điều chỉnh lưới tích hợp là hoạt động chọn các tensor tham số một cách trơn tru (theo tối ưu hóa SGD) mà các bộ lọc của nó sẽ được lấy mẫu cho các số do người dùng xác định. Không giống như các phương pháp xóa bộ lọc/kênh, INN tạo ra các bộ lọc có thể kết hợp một số bộ lọc riêng biệt do hoạt động nội suy.

 # Set trainable gird for each integral layer # Each group should have the same grid # During the sum of continuous signals # We need to sample it using the same set of points for group in model.groups: new_size = 224 if 'operator' in group.operations else 128 group.reset_grid(inn.TrainableGrid1D(new_size)) # Prepare model for tuning of integration grid model.grid_tuning() # Start training train(model, train_data, test_data)

Trong quá trình thử nghiệm, chúng tôi nhận thấy rằng 500 hình ảnh cho kết quả tương tự như một bộ Div2k đầy đủ gồm 4000 hình ảnh.

Hiệu suất

Mô hình INN thu được có thể dễ dàng chuyển đổi thành mô hình riêng biệt và được triển khai trên bất kỳ GPU NVIDIA nào. Chúng tôi cung cấp số khung hình trên giây (FPS) trên RTX A4000 với độ phân giải đầu vào là 64x64. Mô hình nén đạt được tốc độ tăng gần gấp đôi. Để chuyển đổi mô hình INN được cắt bớt sang mô hình rời rạc, có thể sử dụng dòng mã sau:

 model = model.transform_to_discrete() # then model can be compiled, for instance # compilation can add an additional 1.4x speedup for inference model = torch.compile(model, backend='cudagraphs') 

Phần kết luận

Trong bài viết này, chúng tôi đã trình bày tổng quan về bài báo ứng cử viên giải thưởng CVPR2023, “Mạng thần kinh tích hợp”. Nó được áp dụng để cắt tỉa sau đào tạo của mô hình EDSR 4x, đạt được tốc độ tăng gần gấp đôi chỉ với một dòng mã và tinh chỉnh lưới tích hợp trong 1 phút.

Trong các bài viết tiếp theo, chúng tôi sẽ trình bày nhiều ứng dụng hơn của INN và sẽ đề cập đến nhiều chi tiết hơn về việc triển khai mô hình hiệu quả. Giữ nguyên:

Trang web dự án INN

Dự án INN Github

Mã hỗ trợ của bài viết này

Cám ơn vì sự quan tâm của bạn!