关于pytorch的save函数

本文探讨了在PyTorch中保存和加载模型时遇到的问题,特别是`collections.OrderedDict`对象没有'forward'属性的错误。示例代码展示了如何使用`torch.save()`函数分别保存模型的状态字典和整个模型对象。在测试阶段,只保存状态字典会导致模型无法执行前向传播。解决此问题需要完整保存和加载模型实例。
部署运行你感兴趣的模型镜像
 # torch.save(Modul.state_dict(),store_dir)#保存训练模型
        torch.save(Modul,store_dir)#上面的只保存了参数 测试时会'collections.OrderedDict' object has no attribute 'forward'

RT#上面的只保存了参数 测试时会'collections.OrderedDict' object has no attribute 'forward'

注意:input的store_dir需要给出文件名 exp:C:/test/model.pt     错误:C:/test/

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch 2.5

PyTorch
Cuda

PyTorch 是一个开源的 Python 机器学习库,基于 Torch 库,底层由 C++ 实现,应用于人工智能领域,如计算机视觉和自然语言处理

编写PyTorch模型训练函数涉及多个关键步骤,包括模型定义、设备分配、损失函数和优化器的初始化、数据加载、训练循环的构建以及模型性能的评估。以下是完整的实现方法和注意事项。 ### 1. 定义训练函数的基本结构 训练函数通常包括以下几个部分: - 设备分配(CPU/GPU) - 模型初始化 - 数据加载器(DataLoader) - 优化器和损失函数的定义 - 训练循环(epoch和batch级别的迭代) - 模型保存和性能评估 ```python import torch from torch.utils.data import DataLoader import torch.nn as nn import torch.optim as optim def train_model(model, dataloader: DataLoader, num_epochs=10, learning_rate=0.001, device='cuda' if torch.cuda.is_available() else 'cpu'): model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): model.train() running_loss = 0.0 for inputs, labels in dataloader: inputs, labels = inputs.to(device), labels.to(device) # 前向传播 outputs = model(inputs) loss = criterion(outputs, labels) # 反向传播 optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) epoch_loss = running_loss / len(dataloader.dataset) print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}') return model ``` ### 2. 数据加载与预处理 使用`torch.utils.data.DataLoader`来加载数据,支持批量加载和数据增强。可以结合`torchvision.transforms`进行数据预处理: ```python from torchvision import transforms from torchvision.datasets import CIFAR10 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) train_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) ``` ### 3. 模型定义 定义一个简单的卷积神经网络(CNN)作为示例: ```python import torch.nn as nn class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(32 * 8 * 8, 10) def forward(self, x): x = self.pool(torch.relu(self.conv1(x))) x = self.pool(torch.relu(self.conv2(x))) x = x.view(-1, 32 * 8 * 8) x = self.fc1(x) return x ``` ### 4. 模型训练与保存 调用训练函数并传入模型和数据加载器: ```python model = SimpleCNN() trained_model = train_model(model, train_loader, num_epochs=10, learning_rate=0.001) # 保存训练好的模型 torch.save(trained_model.state_dict(), 'trained_model.pth') ``` ### 5. 模型评估 在训练完成后,可以通过测试集评估模型性能: ```python def evaluate_model(model, test_loader, device='cuda' if torch.cuda.is_available() else 'cpu'): model.to(device) model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Test Accuracy: {100 * correct / total:.2f}%') ``` ### 6. 完整代码整合 将上述代码整合为完整的训练流程: ```python import torch from torch.utils.data import DataLoader import torch.nn as nn import torch.optim as optim from torchvision import transforms from torchvision.datasets import CIFAR10 class SimpleCNN(nn.Module): def __init__(self): super(SimpleCNN, self).__init__() self.conv1 = nn.Conv2d(3, 16, 3, padding=1) self.conv2 = nn.Conv2d(16, 32, 3, padding=1) self.pool = nn.MaxPool2d(2, 2) self.fc1 = nn.Linear(32 * 8 * 8, 10) def forward(self, x): x = self.pool(torch.relu(self.conv1(x))) x = self.pool(torch.relu(self.conv2(x))) x = x.view(-1, 32 * 8 * 8) x = self.fc1(x) return x def train_model(model, dataloader: DataLoader, num_epochs=10, learning_rate=0.001, device='cuda' if torch.cuda.is_available() else 'cpu'): model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): model.train() running_loss = 0.0 for inputs, labels in dataloader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() running_loss += loss.item() * inputs.size(0) epoch_loss = running_loss / len(dataloader.dataset) print(f'Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}') return model def evaluate_model(model, test_loader, device='cuda' if torch.cuda.is_available() else 'cpu'): model.to(device) model.eval() correct = 0 total = 0 with torch.no_grad(): for inputs, labels in test_loader: inputs, labels = inputs.to(device), labels.to(device) outputs = model(inputs) _, predicted = torch.max(outputs.data, 1) total += labels.size(0) correct += (predicted == labels).sum().item() print(f'Test Accuracy: {100 * correct / total:.2f}%') # 数据预处理和加载 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) train_dataset = CIFAR10(root='./data', train=True, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_dataset = CIFAR10(root='./data', train=False, download=True, transform=transform) test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False) # 模型训练 model = SimpleCNN() trained_model = train_model(model, train_loader, num_epochs=10, learning_rate=0.001) # 保存模型 torch.save(trained_model.state_dict(), 'trained_model.pth') # 模型评估 evaluate_model(trained_model, test_loader) ``` ### 7. 模型微调(Fine-tuning) 如果使用预训练模型,可以通过加载预训练权重并冻结部分层进行微调[^2]: ```python import torchvision.models as models # 加载预训练的ResNet-50模型 model = models.resnet50(pretrained=True) # 冻结所有参数 for param in model.parameters(): param.requires_grad = False # 替换最后一层 num_ftrs = model.fc.in_features model.fc = nn.Linear(num_ftrs, 10) # 假设新任务有10个类别 ``` ###
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值