• 日期: 2025-04-23
  • 标签: 常规

以下是使用MindSpore框架训练ResNet模型的示例代码，该代码使用了XavierUniform初始化权重，并使用GeneratorDataset加载自定义数据集进行训练。

import numpy as np
import mindspore.dataset as ds
import os
import cv2
import mindspore
import mindspore.nn as nn
from mindspore import Tensor
from mindspore.common.initializer import Normal, XavierUniform
from mindspore import context
from mindspore.train.callback import ModelCheckpoint, CheckpointConfig, LossMonitor, TimeMonitor
from mindspore.train import Model
from mindspore.nn.metrics import Accuracy
from mindspore.ops.operations import TensorAdd
from mindspore import Model # 承载网络结构
from mindspore.nn.metrics import Accuracy # 测试模型用
from mindspore import context
from mindspore.common.initializer import initializer
from mindspore.common.initializer import Normal, XavierUniform
context.set_context(device_target='CPU')

np.random.seed(58)


class BasicBlock(nn.Cell):
    def __init__(self, in_channels, out_channels, stride=1, downsample=None):
        super(BasicBlock, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, has_bias=False, weight_init=XavierUniform())
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU()
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, pad_mode='pad', has_bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.downsample = downsample
        self.add = TensorAdd()

    def construct(self, x):
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out = self.add(out, identity)
        out = self.relu(out)

        return out

class ResNet(nn.Cell):
    def __init__(self, block, layers, num_classes=1000):
        super(ResNet,self).__init__()
        self.inplanes = 64
        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, has_bias=False, weight_init=initializer.XavierUniform())
        self.bn1 = nn.BatchNorm2d(64, eps=1e-05, momentum=0.9, gamma_init=initializer.Constant(1), beta_init=initializer.Zero())
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, pad_mode='same')
        self.layer1 = self._make_layer(block, 64, layers[0])
        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
        self.avgpool = nn.AvgPool2d(7)
        self.flatten = nn.Flatten()
        self.fc = nn.Dense(512 * block.expansion, num_classes, weight_init=initializer.Normal(0.01))

    def _make_layer(self, block, planes, blocks, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.SequentialCell([
                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, has_bias=False, weight_init=initializer.XavierUniform()),
                nn.BatchNorm2d(planes * block.expansion, eps=1e-05, momentum=0.9, gamma_init=initializer.Constant(1), beta_init=initializer.Zero())
            ])
        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes))
        return nn.SequentialCell(layers)

    def construct(self, x):
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        x = self.maxpool(x)

        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)

        x = self.avgpool(x)
        x = self.flatten(x)
        x = self.fc(x)

        return x



class TrainDatasetGenerator:
    def __init__(self, file_path):
        self.file_path = file_path
        self.img_names = os.listdir(file_path)

    def __getitem__(self, index):
        data = cv2.imread(os.path.join(self.file_path, self.img_names[index]))
        label = self.img_names[index].split('_')[0]
        label = int(label)
        data = cv2.cvtColor(data, cv2.COLOR_BGR2RGB)
        data = cv2.resize(data, (224, 224))
        data = data.transpose().astype(np.float32) / 255.
        return data, label

    def __len__(self):
        return len(self.img_names)


def train_resnet():
    context.set_context(mode=context.GRAPH_MODE, device_target='CPU')
    train_dataset_generator = TrainDatasetGenerator('D:/pythonProject7/train1')
    ds_train = ds.GeneratorDataset(train_dataset_generator, ['data', 'label'], shuffle=True)
    ds_train = ds_train.shuffle(buffer_size=10)
    ds_train = ds_train.batch(batch_size=4, drop_remainder=True)
    valid_dataset_generator = TrainDatasetGenerator('D:/pythonProject7/test1')
    ds_valid = ds.GeneratorDataset(valid_dataset_generator, ['data', 'label'], shuffle=True)
    ds_valid = ds_valid.batch(batch_size=4, drop_remainder=True)
    num_classes = len(set(valid_dataset_generator.img_names)) # 获取测试集中的类别数
    network = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=num_classes)
    net_loss = nn.SoftmaxCrossEntropyWithLogits(sparse=True, reduction='mean')
    net_opt = nn.Momentum(network.trainable_params(), learning_rate=0.01, momentum=0.9)
    time_cb = TimeMonitor(data_size=ds_train.get_dataset_size())
    config_ck = CheckpointConfig(save_checkpoint_steps=10, keep_checkpoint_max=10)
    config_ckpt_path = 'D:/pythonProject7/ckpt/'
    ckpoint_cb = ModelCheckpoint(prefix='checkpoint_resnet', directory=config_ckpt_path, config=config_ck)

    model = Model(network, net_loss, net_opt, metrics={'Accuracy': Accuracy()})
    epoch_size = 10
    print('============== Starting Training =============')
    model.train(epoch_size, ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()])

    acc = model.eval(ds_valid)
    print('============== {} ============='.format(acc))
    epoch_size = 10
    print('============== Starting Training =============')
    model.train(epoch_size, ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()])

    acc = model.eval(ds_valid)
    print('============== {} ============='.format(acc))
    epoch_size = 10
    print('============== Starting Training =============')
    model.train(epoch_size, ds_train, callbacks=[time_cb, ckpoint_cb, LossMonitor()])

    acc = model.eval(ds_valid)
    print('============== {} ============='.format(acc))

if __name__ == '__main__':
    train_resnet()

代码说明：

  1. 导入必要的库：numpy, mindspore.dataset, os, cv2, mindspore, mindspore.nn, mindspore.common.initializer, mindspore.train.callback, mindspore.train, mindspore.nn.metrics
  2. 定义BasicBlock类：
    • 使用nn.Conv2d定义卷积层，并使用XavierUniform初始化权重
    • 使用nn.BatchNorm2d定义批归一化层
    • 使用nn.ReLU定义ReLU激活函数
    • 使用TensorAdd进行残差连接
  3. 定义ResNet类：
    • 使用nn.Conv2d, nn.BatchNorm2d, nn.ReLU, nn.MaxPool2d, nn.AvgPool2d, nn.Flatten, nn.Dense构建ResNet网络结构
    • 使用XavierUniform初始化部分卷积层的权重
  4. 定义TrainDatasetGenerator类：
    • 用于加载自定义数据集
    • 使用cv2读取图像数据，并进行预处理
  5. 定义train_resnet函数：
    • 创建训练集和验证集
    • 初始化ResNet模型
    • 定义损失函数net_loss
    • 定义优化器net_opt
    • 使用Model类构建模型
    • 使用ModelCheckpoint保存模型
    • 使用train方法训练模型
    • 使用eval方法评估模型

运行代码：

  1. 将代码保存为.py文件
  2. 确保安装了MindSpore库
  3. 运行代码，并根据代码中的路径设置修改数据集路径

注意：

  • 本代码使用CPU进行训练，如果需要使用GPU进行训练，请将context.set_context(device_target='CPU')改为context.set_context(device_target='GPU')
  • 数据集路径需要根据实际情况进行修改
  • 可以根据需要修改代码中的超参数，例如学习率、训练轮数等
MindSpore ResNet模型训练示例 - 使用XavierUniform初始化权重

原文地址: https://www.cveoy.top/t/topic/mUPc 著作权归作者所有。请勿转载和采集!

免费AI点我,无需注册和登录