Виправлення помилки завантаження моделі PyTorch: _pickle.UnpicklingError: неправильний ключ завантаження, 'x1f'

Вирішення помилки контрольної точки моделі PyTorch: недійсний ключ завантаження

Серверне рішення Python, що використовує правильну обробку файлів і завантаження моделі

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.")

Альтернативне рішення: повторне збереження файлу контрольної точки

Рішення на основі Python для виправлення пошкодженого файлу контрольної точки

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.")

Модульний тест для обох рішень

Модульні тести для перевірки завантаження контрольної точки та цілісності моделі 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()

Часті запитання щодо проблем із завантаженням контрольної точки PyTorch

1. Чому я отримую _pickle.UnpicklingError під час завантаження моєї моделі PyTorch?
2. Ця помилка зазвичай виникає через несумісний або пошкоджений файл контрольної точки. Це також може статися під час використання різних версій PyTorch між збереженням і завантаженням.
3. Як виправити пошкоджений файл контрольної точки PyTorch?
4. Ви можете використовувати zipfile.ZipFile() щоб перевірити, чи є файл ZIP-архівом, або повторно зберегти контрольну точку за допомогою torch.save() після його ремонту.
5. Яка роль state_dict у PyTorch?
6. The state_dict містить ваги та параметри моделі у форматі словника. Завжди зберігайте та завантажуйте state_dict для кращої переносимості.
7. Як я можу завантажити контрольну точку PyTorch на ЦП?
8. Використовуйте map_location='cpu' аргумент в torch.load() щоб переналаштувати тензори з GPU на CPU.
9. Чи можуть контрольні точки PyTorch вийти з ладу через конфлікти версій?
10. Так, старі контрольні точки можуть не завантажуватися в новіших версіях PyTorch. Під час збереження та завантаження рекомендовано використовувати узгоджені версії PyTorch.
11. Як я можу перевірити, чи файл контрольної точки PyTorch не пошкоджено?
12. Спробуйте завантажити файл за допомогою torch.load(). Якщо це не вдається, перевірте файл за допомогою таких інструментів, як zipfile.is_zipfile().
13. Який правильний спосіб збереження та завантаження моделей PyTorch?
14. Завжди економте за допомогою torch.save(model.state_dict()) і завантаження за допомогою model.load_state_dict().
15. Чому моя модель не завантажується на іншому пристрої?
16. Це трапляється, коли тензори зберігаються для GPU, але завантажуються на CPU. використання map_location щоб вирішити це.
17. Як я можу перевірити контрольні точки в різних середовищах?
18. Напишіть модульні тести за допомогою unittest щоб перевірити завантаження моделі на різних налаштуваннях (CPU, GPU, OS).
19. Чи можу я перевірити файли контрольних точок вручну?
20. Так, ви можете змінити розширення на .zip і відкрити його за допомогою zipfile або менеджери архіву для перевірки вмісту.

Подолання помилок завантаження моделі PyTorch

Завантаження контрольних точок PyTorch іноді може викликати помилки через пошкоджені файли або невідповідність версій. Перевіривши формат файлу та використовуючи належні інструменти, як-от zip-файл або перевідображення тензорів, ви можете ефективно відновити навчені моделі та заощадити години повторного навчання.

Розробники повинні дотримуватися найкращих практик, як-от збереження state_dict тільки та перевірка моделей у різних середовищах. Пам’ятайте, що час, витрачений на вирішення цих проблем, гарантує, що ваші моделі залишаться функціональними, портативними та сумісними з будь-якою системою розгортання. 🚀