模型保存与加载#

import torch

# 保存整个模型
torch.save(model, 'model.pth')

# 加载整个模型
model = torch.load('model.pth')
model.eval()

# 保存状态字典
torch.save(model.state_dict(), 'model_state.pth')

# 加载状态字典
model = MyModel()  # 需要先创建模型实例
model.load_state_dict(torch.load('model_state.pth'))
model.eval()

# 保存检查点
checkpoint = {
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
    'accuracy': accuracy
}
torch.save(checkpoint, 'checkpoint.pth')

# 加载检查点
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

def save_checkpoint(state, is_best, filename='checkpoint.pth', best_filename='best_model.pth'):
    torch.save(state, filename)
    if is_best:
        torch.save(state, best_filename)

def load_checkpoint(checkpoint_path, model, optimizer=None):
    checkpoint = torch.load(checkpoint_path)
    model.load_state_dict(checkpoint['model_state_dict'])
    
    if optimizer is not None:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    epoch = checkpoint['epoch']
    loss = checkpoint.get('loss', None)
    accuracy = checkpoint.get('accuracy', None)
    
    return epoch, loss, accuracy

# 使用示例
best_acc = 0.0
for epoch in range(num_epochs):
    # ... 训练 ...
    
    # 保存检查点
    checkpoint = {
        'epoch': epoch + 1,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict(),
        'loss': val_loss,
        'accuracy': val_acc,
    }
    
    is_best = val_acc > best_acc
    if is_best:
        best_acc = val_acc
    
    save_checkpoint(checkpoint, is_best)

# 保存时指定 map_location
torch.save(model.state_dict(), 'model.pth')

# 加载到 CPU（即使模型在 GPU 上训练）
model.load_state_dict(torch.load('model.pth', map_location='cpu'))

# 加载到特定 GPU
model.load_state_dict(torch.load('model.pth', map_location='cuda:0'))

# 自动选择设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load('model.pth', map_location=device))

# 保存 DataParallel 模型
if isinstance(model, nn.DataParallel):
    model_state_dict = model.module.state_dict()
else:
    model_state_dict = model.state_dict()
torch.save(model_state_dict, 'model.pth')

# 加载到单 GPU
model = MyModel()
model.load_state_dict(torch.load('model.pth', map_location='cpu'))
model = model.to(device)

# 只加载匹配的层
pretrained_dict = torch.load('pretrained.pth')
model_dict = model.state_dict()

# 过滤掉不匹配的键
pretrained_dict = {k: v for k, v in pretrained_dict.items() 
                   if k in model_dict and model_dict[k].shape == v.shape}

# 更新模型字典
model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

# 加载预训练权重
pretrained = torch.load('pretrained.pth')

# 创建新模型（结构可能不同）
model = MyNewModel()

# 获取模型状态字典
model_dict = model.state_dict()

# 匹配并加载权重
pretrained_dict = {k: v for k, v in pretrained.items() 
                   if k in model_dict and v.size() == model_dict[k].size()}

model_dict.update(pretrained_dict)
model.load_state_dict(model_dict)

# 打印未加载的层
missing_keys = set(model_dict.keys()) - set(pretrained_dict.keys())
print(f"未加载的层: {missing_keys}")

import torch.onnx

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出为 ONNX
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
)

# 方法1：Tracing
dummy_input = torch.randn(1, 3, 224, 224)
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save('model_traced.pt')

# 方法2：Scripting（需要模型代码支持）
scripted_model = torch.jit.script(model)
scripted_model.save('model_scripted.pt')

# 加载
loaded_model = torch.jit.load('model_traced.pt')

import torch.quantization

# 量化模型
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

# 保存量化模型
torch.save(quantized_model.state_dict(), 'quantized_model.pth')

model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model_prepared = torch.quantization.prepare(model)
# ... 用校准数据运行模型 ...
model_quantized = torch.quantization.convert(model_prepared)

best_val_acc = 0.0
best_model_state = None

for epoch in range(num_epochs):
    # ... 训练和验证 ...
    
    if val_acc > best_val_acc:
        best_val_acc = val_acc
        best_model_state = model.state_dict().copy()

# 保存最佳模型
torch.save(best_model_state, 'best_model.pth')

save_interval = 10  # 每10个epoch保存一次

for epoch in range(num_epochs):
    # ... 训练 ...
    
    if (epoch + 1) % save_interval == 0:
        checkpoint = {
            'epoch': epoch + 1,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
        }
        torch.save(checkpoint, f'checkpoint_epoch_{epoch+1}.pth')

config = {
    'model_name': 'ResNet18',
    'num_classes': 10,
    'input_size': (3, 224, 224),
    'batch_size': 32,
    'learning_rate': 0.001,
    'num_epochs': 100,
}

checkpoint = {
    'config': config,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'epoch': epoch,
    'loss': loss,
}

torch.save(checkpoint, 'checkpoint.pth')

# 加载时
checkpoint = torch.load('checkpoint.pth')
config = checkpoint['config']
model = create_model_from_config(config)
model.load_state_dict(checkpoint['model_state_dict'])

import datetime

timestamp = datetime.datetime.now().strftime('%Y%m%d_%H%M%S')
checkpoint_path = f'checkpoints/checkpoint_{timestamp}.pth'
torch.save(checkpoint, checkpoint_path)

def load_state_dict_with_missing_keys(model, checkpoint_path):
    checkpoint = torch.load(checkpoint_path)
    model_dict = model.state_dict()
    
    # 检查缺失的键
    missing_keys = set(model_dict.keys()) - set(checkpoint.keys())
    unexpected_keys = set(checkpoint.keys()) - set(model_dict.keys())
    
    if missing_keys:
        print(f"警告：缺失的键: {missing_keys}")
    if unexpected_keys:
        print(f"警告：意外的键: {unexpected_keys}")
    
    # 加载匹配的键
    model_dict.update({k: v for k, v in checkpoint.items() if k in model_dict})
    model.load_state_dict(model_dict, strict=False)
    
    return missing_keys, unexpected_keys

def load_with_shape_check(model, checkpoint_path):
    checkpoint = torch.load(checkpoint_path)
    model_dict = model.state_dict()
    
    pretrained_dict = {}
    for k, v in checkpoint.items():
        if k in model_dict:
            if v.shape == model_dict[k].shape:
                pretrained_dict[k] = v
            else:
                print(f"跳过 {k}：形状不匹配 {v.shape} vs {model_dict[k].shape}")
    
    model_dict.update(pretrained_dict)
    model.load_state_dict(model_dict)

def save_model(model, optimizer, scheduler, epoch, loss, accuracy, filepath, is_best=False):
    """保存完整的训练状态"""
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'scheduler_state_dict': scheduler.state_dict() if scheduler else None,
        'loss': loss,
        'accuracy': accuracy,
    }
    torch.save(checkpoint, filepath)
    
    if is_best:
        best_path = filepath.replace('.pth', '_best.pth')
        torch.save(checkpoint, best_path)

def load_model(model, optimizer, scheduler, filepath, device='cpu'):
    """加载完整的训练状态"""
    checkpoint = torch.load(filepath, map_location=device)
    
    model.load_state_dict(checkpoint['model_state_dict'])
    model.to(device)
    
    if optimizer and 'optimizer_state_dict' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    if scheduler and checkpoint.get('scheduler_state_dict'):
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
    
    epoch = checkpoint.get('epoch', 0)
    loss = checkpoint.get('loss', None)
    accuracy = checkpoint.get('accuracy', None)
    
    return epoch, loss, accuracy