模型训练往往需要较高的配置，为了满足友友们的好奇心，这里我们不要内存，不要gpu，用最简单的方式，让大家感受一下什么是模型训练。基于你的硬件配置，我们可以设计一个完全在CPU上运行的简易模型训练方案。以下是具体步骤：

环境准备

这里以mac为例，其他系统原理类似，也可不使用miniconda，本文主要集中在训练代码和推理代码上。

安装Miniconda（推荐）

# 下载Miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh
# 安装
bash Miniconda3-latest-MacOSX-arm64.sh

创建虚拟环境

conda create -n tinyai python=3.9
conda activate tinyai

安装PyTorch

# 安装pytorch，也可通过官网选择合适的安装语句
pip install torch torchvision torchaudio

超简易模型训练方案

纯CPU训练微型文本模型

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

# 超简单数据集
class TextDataset(Dataset):
    def __init__(self):
        self.data = ["hello world", "deep learning", "apple silicon", "metal acceleration"]
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        text = self.data[idx]
        # 简单字符级编码
        x = [ord(c) for c in text[:-1]]
        y = [ord(c) for c in text[1:]]
        return torch.tensor(x), torch.tensor(y)

# 超简单模型
class TinyLM(nn.Module):
    def __init__(self, vocab_size=128):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, 32)
        self.rnn = nn.RNN(32, 64, batch_first=True)
        self.fc = nn.Linear(64, vocab_size)
    
    def forward(self, x):
        x = self.embed(x)
        out, _ = self.rnn(x)
        return self.fc(out)
def custom_collate_fn(batch):
    # batch是包含多个(__getitem__返回结果)的列表
    x_batch, y_batch = zip(*batch)
    
    # 找到本批次中的最大长度
    max_len = max(len(x) for x in x_batch)
    
    # 填充每个样本
    x_padded = torch.stack([
        torch.cat([x, torch.zeros(max_len - len(x), dtype=torch.long)]) 
        for x in x_batch
    ])
    
    y_padded = torch.stack([
        torch.cat([y, torch.zeros(max_len - len(y), dtype=torch.long)]) 
        for y in y_batch
    ])
    
    return x_padded, y_padded



# 训练设置
dataset = TextDataset()
# loader = DataLoader(dataset, batch_size=2)
# 然后修改DataLoader
loader = DataLoader(dataset, batch_size=2, collate_fn=custom_collate_fn)
model = TinyLM()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练循环
for epoch in range(10):
    for x, y in loader:
        optimizer.zero_grad()
        output = model(x)
        loss = criterion(output.view(-1, 128), y.view(-1))
        loss.backward()
        optimizer.step()
    print(f"Epoch {epoch}, Loss: {loss.item():.4f}")

# 保存模型和tokenizer（虽然我们用的是简单编码）
torch.save(model.state_dict(), 'tinylm.pth')

# 同时保存词汇表信息（这里只是示例，实际字符编码是固定的）
import pickle
with open('char_vocab.pkl', 'wb') as f:
    pickle.dump({'vocab_size': 128}, f)  # ASCII字符范围

模型推理

创建一个新的Python文件inference.py：

import torch
import torch.nn as nn

class TinyLM(nn.Module):
    def __init__(self, vocab_size=128):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, 32)
        self.rnn = nn.RNN(32, 64, batch_first=True)
        self.fc = nn.Linear(64, vocab_size)
    
    def forward(self, x):
        x = self.embed(x)
        out, _ = self.rnn(x)
        return self.fc(out)

# 加载模型
model = TinyLM()
model.load_state_dict(torch.load('tinylm.pth'))
model.eval()  # 设置为评估模式

# 简单的字符编码函数
def text_to_tensor(text):
    return torch.tensor([[ord(c) for c in text]])

# 推理函数
def generate_text(start_str, length=10):
    input_seq = text_to_tensor(start_str)
    hidden = None
    
    for _ in range(length):
        with torch.no_grad():  # 禁用梯度计算
            output = model(input_seq)
            # 获取最后一个字符的预测
            last_char_logits = output[0, -1, :]
            # 选择概率最高的字符
            predicted_char = torch.argmax(last_char_logits).item()
            # 添加到输入序列中
            input_seq = torch.cat([
                input_seq, 
                torch.tensor([[predicted_char]])
            ], dim=1)
    
    # 将数字转换回字符
    generated_text = ''.join([chr(c) for c in input_seq[0].tolist()])
    return generated_text

# 使用示例
if __name__ == "__main__":
    while True:
        seed = input("输入起始字符串(或输入q退出): ")
        if seed.lower() == 'q':
            break
        generated = generate_text(seed, length=20)
        print(f"生成结果: {generated}")

运行推理示例

输入起始字符串(或输入q退出): hello
生成结果: hello world deep lear

结果分析

生成的文本无意义

主要原因是模型太小或训练不足，后续的解决方案是增加训练epoch或扩大模型，当然本文的目的就是让大家熟悉一下基本的模型训练和推理流程。