Sửa lỗi tải mô hình PyTorch: _pickle.UnpicklingError: khóa tải không hợp lệ, 'x1f'

Giải quyết lỗi điểm kiểm tra mô hình PyTorch: Khóa tải không hợp lệ

Giải pháp phụ trợ Python sử dụng xử lý tệp và tải mô hình thích hợp

import os
import torch
import numpy as np
import timm
import zipfile
import io
# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device being used:', device)
# Correct method to load a corrupted or zipped model checkpoint
mname = os.path.join('./CDF2_0.pth')
try:
    # Attempt to open as a zip if initial loading fails
    if zipfile.is_zipfile(mname):
        with zipfile.ZipFile(mname) as archive:
            for file in archive.namelist():
                with archive.open(file) as f:
                    buffer = io.BytesIO(f.read())
                    checkpoints = torch.load(buffer, map_location=device)
    else:
        checkpoints = torch.load(mname, map_location=device)
    print("Checkpoint loaded successfully.")
except Exception as e:
    print("Error loading the checkpoint file:", e)
# Model creation and state_dict loading
model = timm.create_model('legacy_xception', pretrained=True, num_classes=2).to(device)
if 'state_dict' in checkpoints:
    model.load_state_dict(checkpoints['state_dict'])
else:
    model.load_state_dict(checkpoints)
model.eval()
print("Model loaded and ready for inference.")

Giải pháp thay thế: Lưu lại tệp điểm kiểm tra

Giải pháp dựa trên Python để sửa tệp điểm kiểm tra bị hỏng

import os
import torch
# Device setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Device being used:', device)
# Original and corrected file paths
original_file = './CDF2_0.pth'
corrected_file = './fixed_CDF2_0.pth'
try:
    # Load and re-save the checkpoint
    checkpoints = torch.load(original_file, map_location=device)
    torch.save(checkpoints, corrected_file)
    print("Checkpoint file re-saved successfully.")
except Exception as e:
    print("Failed to fix checkpoint file:", e)
# Verify loading from the corrected file
checkpoints_fixed = torch.load(corrected_file, map_location=device)
print("Verified: Corrected checkpoint loaded.")

Kiểm tra đơn vị cho cả hai giải pháp

Kiểm tra đơn vị để xác thực việc tải điểm kiểm tra và tính toàn vẹn của mô hình state_dict

import torch
import unittest
import os
import timm
class TestCheckpointLoading(unittest.TestCase):
    def setUp(self):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model_path = './fixed_CDF2_0.pth'
        self.model = timm.create_model('legacy_xception', pretrained=True, num_classes=2).to(self.device)
    def test_checkpoint_loading(self):
        try:
            checkpoints = torch.load(self.model_path, map_location=self.device)
            if 'state_dict' in checkpoints:
                self.model.load_state_dict(checkpoints['state_dict'])
            else:
                self.model.load_state_dict(checkpoints)
            self.model.eval()
            self.assertTrue(True)
            print("Checkpoint loaded successfully in unit test.")
        except Exception as e:
            self.fail(f"Checkpoint loading failed with error: {e}")
if __name__ == '__main__':
    unittest.main()

Câu hỏi thường gặp về vấn đề tải điểm kiểm tra PyTorch

1. Tại sao tôi lại nhận được _pickle.UnpicklingError khi tải mô hình PyTorch của tôi?
2. Lỗi này thường xảy ra do file checkpoint không tương thích hoặc bị hỏng. Nó cũng có thể xảy ra khi sử dụng các phiên bản PyTorch khác nhau giữa quá trình lưu và tải.
3. Làm cách nào để sửa tệp điểm kiểm tra PyTorch bị hỏng?
4. Bạn có thể sử dụng zipfile.ZipFile() để kiểm tra xem tệp có phải là kho lưu trữ ZIP hay lưu lại điểm kiểm tra bằng torch.save() sau khi sửa chữa nó.
5. Vai trò của state_dict trong PyTorch?
6. các state_dict chứa các trọng số và tham số của mô hình ở định dạng từ điển. Luôn lưu và tải state_dict để có tính di động tốt hơn.
7. Làm cách nào tôi có thể tải điểm kiểm tra PyTorch trên CPU?
8. Sử dụng map_location='cpu' tranh luận trong torch.load() để ánh xạ lại các tensor từ GPU sang CPU.
9. Điểm kiểm tra PyTorch có thể không thành công do xung đột phiên bản không?
10. Có, các điểm kiểm tra cũ hơn có thể không tải được trong các phiên bản PyTorch mới hơn. Bạn nên sử dụng các phiên bản PyTorch nhất quán khi lưu và tải.
11. Làm cách nào để kiểm tra xem tệp điểm kiểm tra PyTorch có bị hỏng không?
12. Hãy thử tải tập tin bằng cách sử dụng torch.load(). Nếu thất bại, hãy kiểm tra tệp bằng các công cụ như zipfile.is_zipfile().
13. Cách chính xác để lưu và tải các mô hình PyTorch là gì?
14. Luôn lưu bằng cách sử dụng torch.save(model.state_dict()) và tải bằng cách sử dụng model.load_state_dict().
15. Tại sao mô hình của tôi không tải được trên một thiết bị khác?
16. Điều này xảy ra khi các tensor được lưu cho GPU nhưng được tải trên CPU. Sử dụng map_location để giải quyết việc này.
17. Làm cách nào tôi có thể xác thực các điểm kiểm tra trên các môi trường?
18. Viết bài kiểm tra đơn vị bằng cách sử dụng unittest để kiểm tra tải mô hình trên các thiết lập khác nhau (CPU, GPU, HĐH).
19. Tôi có thể kiểm tra các tập tin điểm kiểm tra theo cách thủ công không?
20. Có, bạn có thể thay đổi phần mở rộng thành .zip và mở nó bằng zipfile hoặc người quản lý lưu trữ để kiểm tra nội dung.

Khắc phục lỗi tải mô hình PyTorch

Việc tải điểm kiểm tra PyTorch đôi khi có thể gây ra lỗi do tệp bị hỏng hoặc phiên bản không khớp. Bằng cách xác minh định dạng tệp và sử dụng các công cụ thích hợp như tập tin zip hoặc ánh xạ lại các tensor, bạn có thể khôi phục các mô hình đã đào tạo của mình một cách hiệu quả và tiết kiệm hàng giờ đào tạo lại.

Nhà phát triển nên làm theo các phương pháp hay nhất như lưu trạng thái_dict duy nhất và xác nhận các mô hình trên các môi trường. Hãy nhớ rằng, thời gian dành để giải quyết những vấn đề này đảm bảo mô hình của bạn vẫn hoạt động, di động và tương thích với mọi hệ thống triển khai. 🚀