pytorch中神经网络模型的初始化-保存-加载

博主打算训练自己的模型以获得较好的training_model,包含模型初始化。介绍了三种模型初始化方式,第一种需先定义初始化函数再调用,且要实例化网络模型;第三种是在utils文件中定义initialize_weights函数,还给出初始化位置照片。
部署运行你感兴趣的模型镜像

https://blog.youkuaiyun.com/remanented/article/details/89161297

一、打算开始训练自己的模型,希望能够得到较好的training_model,包括了对模型的初始化

第一种

from torch.nn import init
#define the initial function to init the layer's parameters for the network
def weigth_init(m):
    if isinstance(m, nn.Conv2d):
        init.xavier_uniform_(m.weight.data)
        init.constant_(m.bias.data,0.1)
    elif isinstance(m, nn.BatchNorm2d):
        m.weight.data.fill_(1)
        m.bias.data.zero_()
    elif isinstance(m, nn.Linear):
        m.weight.data.normal_(0,0.01)
        m.bias.data.zero_()

首先定义了一个初始化函数,接着进行调用就ok了,不过要先把网络模型实例化: 

    #Define Network
    model = Net(args.input_channel,args.output_channel)
    model.apply(weigth_init)

第二种

def initNetParams(net):
    '''Init net parameters.'''
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            init.xavier_uniform(m.weight)
            if m.bias:
                init.constant(m.bias, 0)
        elif isinstance(m, nn.BatchNorm2d):
            init.constant(m.weight, 1)
            init.constant(m.bias, 0)
        elif isinstance(m, nn.Linear):
            init.normal(m.weight, std=1e-3)
            if m.bias:
                init.constant(m.bias, 0)
 
initNetParams(net)

第三种(我自己使用的)

在utils文件中定义initialize_weights函数

def initialize_weights(net):
    for m in net.modules():
        if isinstance(m, nn.Conv2d):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()
        elif isinstance(m, nn.ConvTranspose2d):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()
        elif isinstance(m, nn.Linear):
            m.weight.data.normal_(0, 0.02)
            m.bias.data.zero_()

 此处是照片,告诉大家初始化的位置,代码在下面

 

import torch
import utils
# 定义网络结构
class CNNnet(torch.nn.Module):
    def __init__(self):
        super(CNNnet,self).__init__()
        self.conv1 = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=1,
                            out_channels=16,
                            kernel_size=3,
                            stride=2,
                            padding=1),
            torch.nn.BatchNorm2d(16),
            torch.nn.ReLU()
        )
        self.conv2 = torch.nn.Sequential(
            torch.nn.Conv2d(16,32,3,2,1),
            torch.nn.BatchNorm2d(32),
            torch.nn.ReLU()
        )
        self.conv3 = torch.nn.Sequential(
            torch.nn.Conv2d(32,64,3,2,1),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU()
        )
        self.conv4 = torch.nn.Sequential(
            torch.nn.Conv2d(64,64,2,2,0),
            torch.nn.BatchNorm2d(64),
            torch.nn.ReLU()
        )
        self.mlp1 = torch.nn.Linear(2*2*64,100)
        self.mlp2 = torch.nn.Linear(100,10)

        utils.initialize_weights(self)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.mlp1(x.view(x.size(0),-1))
        x = self.mlp2(x)
        return x

以上两个文件结合就可以

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

Yolo-v8.3

Yolo-v8.3

Yolo

YOLO(You Only Look Once)是一种流行的物体检测和图像分割模型,由华盛顿大学的Joseph Redmon 和Ali Farhadi 开发。 YOLO 于2015 年推出,因其高速和高精度而广受欢迎

<think>我们正在使用PyTorch保存训练好的神经网络模型。根据引用[1]和引用[4],PyTorch提供了多种保存模型的方法。下面将详细介绍几种主要方法,并给出示例代码。 ### 方法1: 保存整个模型 这种方法会保存模型的架构、权重以及优化器状态等所有信息。使用`torch.save()`函数直接保存模型对象。 ```python import torch import torch.nn as nn # 假设我们有一个训练好的模型 class MyModel(nn.Module): def __init__(self): super(MyModel, self).__init__() self.fc1 = nn.Linear(28*28, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = torch.flatten(x, 1) x = torch.relu(self.fc1(x)) x = self.fc2(x) return x model = MyModel() # ... 训练过程 ... # 保存整个模型 torch.save(model, 'model.pth') ``` 加载整个模型时,直接使用`torch.load()`: ```python model = torch.load('model.pth') model.eval() # 设置为评估模式 ``` **注意**:这种方法在加载时依赖于原始模型类的定义。如果保存后修改了模型类,加载可能会出错。 ### 方法2:保存模型的状态字典(推荐) 更常用的方法是只保存模型的状态字典(`state_dict`),它包含模型的所有权重参数。这种方法更灵活,且不依赖于模型类的具体结构。 ```python # 保存模型的状态字典 torch.save(model.state_dict(), 'model_weights.pth') ``` 加载时,需要先创建模型实例,然后加载状态字典: ```python # 创建相同结构的模型 loaded_model = MyModel() # 加载保存的权重 loaded_model.load_state_dict(torch.load('model_weights.pth')) loaded_model.eval() # 设置为评估模式 ``` ### 方法3: 保存检查点(包含额外信息) 在训练过程中,我们通常希望保存检查点(checkpoint),以便在中断后恢复训练。检查点除了模型权重,还可以包含优化器状态、当前epoch、损失值等。 ```python # 假设我们还有优化器 optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 定义检查点字典 checkpoint = { 'epoch': 10, # 当前训练轮次 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': 0.05, # 当前损失值 } # 保存检查点 torch.save(checkpoint, 'checkpoint.pth') ``` 加载检查点并恢复训练: ```python # 创建模型和优化器 model = MyModel() optimizer = torch.optim.SGD(model.parameters(), lr=0.01) # 加载检查点 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'] # 恢复训练 model.train() # 设置为训练模式 # ... 继续训练 ... ``` ### 关键注意事项 1. **模型模式**:在保存前,确保模型处于正确的模式(训练模式或评估模式)。加载后,根据用途调用`model.train()`或`model.eval()`。 2. **设备映射**:当在GPU上训练而需要在CPU上加载时,可以使用`map_location`参数: ```python model.load_state_dict(torch.load('model_weights.pth', map_location=torch.device('cpu'))) ``` 3. **文件格式**:PyTorch推荐使用`.pth`或`.pt`作为模型文件扩展名,但实际可以是任意扩展名[^1]。 4. **兼容性**:不同版本的PyTorch保存模型可能存在兼容性问题,尽量保持环境一致。 ### 完整示例流程 以下是一个完整的训练并保存模型的示例(基于引用[3]的Fashion-MNIST数据集): ```python import torch import torch.nn as nn import torch.optim as optim from torchvision import datasets, transforms # 定义神经网络(同引用[3]) class QYNN(nn.Module): def __init__(self): super().__init__() self.fc1 = nn.Linear(28 * 28, 128) self.fc2 = nn.Linear(128, 10) def forward(self, x): x = torch.flatten(x, start_dim=1) x = torch.relu(self.fc1(x)) x = self.fc2(x) return x # 加载数据集 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset = datasets.FashionMNIST('./data', download=True, train=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True) # 初始化模型和优化器 model = QYNN() optimizer = optim.SGD(model.parameters(), lr=0.01) criterion = nn.CrossEntropyLoss() # 训练模型 num_epochs = 5 for epoch in range(num_epochs): for images, labels in train_loader: optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() print(f'Epoch {epoch+1}, Loss: {loss.item():.4f}') # 保存模型的状态字典(推荐方法) torch.save(model.state_dict(), 'fashion_mnist_model.pth') # 保存整个模型(备选) torch.save(model, 'fashion_mnist_full_model.pth') ``` ### 总结 - **推荐方法**:保存状态字典(`state_dict`)是最灵活和常用的方式。 - **检查点**:在训练过程中保存检查点,以便恢复训练。 - **部署**:保存模型可以轻松加载用于推理(预测)或迁移学习。 §§相关问题§§ 1. 如何在PyTorch加载保存模型进行预测? 2. 训练过程中如何定期保存检查点? 3. 如何将PyTorch模型转换为ONNX格式以用于其他框架? 4. 在加载模型时遇到“Missing key(s) in state_dict”错误如何解决?</think>### 如何在PyTorch保存训练好的神经网络模型PyTorch保存训练好的模型是深度学习工作流的关键环节,主要分为三种方法,每种方法适用于不同场景。以下是详细步骤和代码示例: --- #### 方法1: 保存整个模型(最简单) 保存模型架构、权重和优化器状态,适合快速部署: ```python import torch # 定义模型(示例引用[3]) class QYNN(torch.nn.Module): def __init__(self): super().__init__() self.fc1 = torch.nn.Linear(28*28, 128) self.fc2 = torch.nn.Linear(128, 10) def forward(self, x): x = torch.flatten(x, 1) x = torch.relu(self.fc1(x)) return self.fc2(x) # 训练模型... model = QYNN() # 保存整个模型(含结构和参数) torch.save(model, 'full_model.pth') # 文件扩展名通常为.pth或.pt ``` **加载方式**: ```python loaded_model = torch.load('full_model.pth') loaded_model.eval() # 切换为评估模式 ``` **优点**:一键保存/加载 **缺点**:文件较大,加载依赖原始类定义[^1] --- #### 方法2:保存模型参数(推荐) 仅保存权重参数,最轻量且灵活的方式: ```python # 保存状态字典(state_dict) torch.save(model.state_dict(), 'model_weights.pth') ``` **加载方式**: ```python # 需先实例化相同结构的模型 new_model = QYNN() new_model.load_state_dict(torch.load('model_weights.pth')) new_model.eval() ``` **优点**: - 文件小(仅参数) - 不依赖原始类代码 - 支持跨设备加载(如GPU→CPU) ```python # 示例:GPU模型加载到CPU new_model.load_state_dict(torch.load('model_weights.pth', map_location=torch.device('cpu'))) ``` [^4] --- #### 方法3: 保存检查点(训练恢复) 保存训练状态(含优化器、epoch等),用于中断后恢复训练: ```python checkpoint = { 'epoch': 10, # 当前训练轮次 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': 0.023, # 当前损失值 } torch.save(checkpoint, 'checkpoint.pth') ``` **恢复训练**: ```python checkpoint = torch.load('checkpoint.pth') model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) current_epoch = checkpoint['epoch'] ``` **适用场景**:长时间训练任务、模型微调[^1] --- ### 关键注意事项 1. **模式切换**:预测前务必调用 `model.eval()` 关闭Dropout/BatchNorm 2. **设备兼容**:使用 `map_location` 参数处理GPU/CPU转换 3. **文件格式**: - `.pth`:PyTorch标准格式 - `.pt`:替代扩展名 - `.tar`:检查点常用格式 4. **版本控制**:不同PyTorch版本的模型可能不兼容 > **数学原理补充**:模型保存本质是存储权重矩阵 $W$ 和偏置向量 $b$。对于全连接层,前向传播计算为: > $$\mathbf{y} = \sigma(W\mathbf{x} + b)$$ > 其中 $\sigma$ 是激活函数(如ReLU)。保存的state_dict即包含所有 $W$ 和 $b$ 的值[^2]。 --- ### 完整训练到保存流程 基于Fashion-MNIST示例(引用[3]): ```python # 1. 数据准备 transform = torchvision.transforms.Compose([ torchvision.transforms.ToTensor(), torchvision.transforms.Normalize((0.5,), (0.5,)) ]) train_set = torchvision.datasets.FashionMNIST('./data', download=True, transform=transform) train_loader = torch.utils.data.DataLoader(train_set, batch_size=32) # 2. 模型训练 model = QYNN() optimizer = torch.optim.Adam(model.parameters()) criterion = torch.nn.CrossEntropyLoss() for epoch in range(10): for images, labels in train_loader: outputs = model(images) loss = criterion(outputs, labels) optimizer.zero_grad() loss.backward() optimizer.step() # 3. 保存模型参数(推荐) torch.save(model.state_dict(), 'fashion_mnist_model.pth') ```
评论 2
成就一亿技术人!
拼手气红包6.0元
还能输入1000个字符
 
红包 添加红包
表情包 插入表情
 条评论被折叠 查看
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值