DeepSpeed运行一个简单神经网络

实现代码和注释如下

#导入需要的包
from accelerate import Accelerator, DeepSpeedPlugin
import torch
from torch.utils.data import DataLoader, TensorDataset
from tqdm import tqdm
from time import time

#定义一个简单的神经网络
class SimpleNet(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleNet, self).__init__() #继承基类的初始化方式
        self.fc1 = torch.nn.Linear(input_size, hidden_size) 
        #定义全连接层，将输入层映射到隐藏层
        self.fc2 = torch.nn.Linear(hidden_size, output_size)
        #定义全连接层，将隐藏层映射到输出层
    
    def forward(self,x): #定义前向函数
        x = torch.relu(self.fc1(x)) #使用relu激活函数
        x = self.fc2(x)
        return x


if __name__ == "__main__":
    #定义相关的参数大小，
    input_size = 10 #输入层，
    hidden_size = 16 #隐藏层，
    output_size = 8 #输出层，
    batch_size = 64 #batch，
    data_size = 10000 #数据量

    input_data = torch.randn(data_size, input_size) #随机初始化输入数据
    labels = torch.randn(data_size, output_size) #随机初始化相应的标签
    
    #使用TensorDataset将输入数据和标签转化为tensor张量并存入dataset
    dataset = TensorDataset(input_data, labels) 
    #使用DataLoader将存储张量的dataset按照batch_size的大小分割，并且放入dataloader
    #shuffle参数表示分割时是否将这些数据随机打乱
    dataloader = DataLoader(dataset, batch_size = batch_size, shuffle = True)
    
    #初始化模型
    model = SimpleNet(input_size, hidden_size, output_size)
    
    #定义一个deepspeed的种子，参数设置为stage = 2，clipping = 1.0
    deepspeed_plugin = DeepSpeedPlugin(zero_stage = 2, gradient_clipping = 1.0)
    #将定义好的deepspeed种子通过Accelerator传入accelerator参数，方便之后调用
    accelerator = Accelerator(deepspeed_plugin = deepspeed_plugin)
    
    #定义优化器，这里使用的SGD，将模型参数model.parameters()传入，学习率为1e-4
    optimizer = torch.optim.SGD(model.parameters(), lr = 1e-4)
    #定义损失函数为交叉熵损失函数
    crition = torch.nn.CrossEntropyLoss()
    
    #使用Accelerator时必要的一步，使用accelerator.prepare将和模型有关的参数传入，
    #accelerator会将这些参数重新返回，再调用时便会触发deepspeed
    model, optimizer, dataloader = accelerator.prepare(model, optimizer, dataloader)
    
    #初始化一个time()来记录训练开始时的时间
    start_time = time()
    #训练100个Epoch
    for epoch in range(100):
        #模型训练初始化
        model.train()
        #从dataloader中取出一个batch的数据
        #tqdm会根据目前取到第几个batch给训练添加进度条
        for batch in tqdm(dataloader):
            #将batch中的数据分开为输入数据和对应的标签
            input, label = batch
            #输入数据到模型中，并得到输出
            output = model(input)
            #计算输出结果和实际结果之间的损失
            loss = crition(output, label)
            
            #每次迭代将模型中所有可训练参数的梯度清零，防止出现梯度累加影响之后的计算
            optimizer.zero_grad()
            #使用accelerator计算loss的梯度，进行反向传播
            accelerator.backward(loss)
            #更新模型参数
            optimizer.step()
        #打印目前Epoch和损失值
        print(f"Epoch: {epoch}, Loss: {loss.item()}")
    #初始化一个time()来记录训练结束时的时间
    end_time = time()
    #计算训练总时长
    total_time = end_time - start_time
    #打印总时长
    print(f"Total_time: {total_time}")
    #使用accelerator保存模型，命名为model.pth
    accelerator.save(model, "model.pth")
    

测试结果

在解释器中运行，等待便可以得到如下输出结果