如何扩展picoGPT功能：添加温度采样和top-p策略的完整教程

【免费下载链接】picoGPT An unnecessarily tiny implementation of GPT-2 in NumPy. 项目地址: https://gitcode.com/gh_mirrors/pi/picoGPT

picoGPT是一个极简的GPT-2实现项目，使用纯NumPy编写，代码精简但功能完整。本文将为你展示如何为picoGPT添加温度采样和top-p策略，让这个微小的GPT实现具备更强大的文本生成能力。通过本教程，你将学会如何扩展picoGPT的采样策略，从简单的贪婪采样升级到更灵活的随机采样方法。

为什么需要扩展picoGPT的采样策略？

在阅读gpt2.py和gpt2_pico.py文件时，你会发现当前的picoGPT只支持贪婪采样（greedy sampling）。这意味着每次生成下一个token时，它总是选择概率最高的那个token。虽然这种方法简单高效，但会导致生成的文本过于确定和重复。

温度采样（temperature sampling）和top-p采样（nucleus sampling）是两种广泛使用的文本生成策略，能够：

1. 增加文本多样性：通过引入随机性，生成更具创造性的内容
2. 控制生成质量：调整温度参数可以平衡生成文本的保守性和创造性
3. 避免重复循环：防止模型陷入重复的模式

理解picoGPT的当前采样实现

首先，让我们看一下picoGPT当前的采样代码。在gpt2.py的第90-92行：

logits = gpt2(inputs, **params, n_head=n_head)  # model forward pass
next_id = np.argmax(logits[-1])  # greedy sampling
inputs.append(int(next_id))  # append prediction to input

这段代码使用np.argmax()函数从logits中直接选择概率最高的token。这种贪婪采样方法虽然简单，但缺乏灵活性。

添加温度采样功能

温度采样通过调整logits的分布来控制生成文本的随机性。温度值（temperature）越高，生成的文本越随机；温度值越低，生成的文本越保守。

实现温度采样函数

我们需要在gpt2.py中添加一个新的采样函数。首先，在文件顶部添加必要的导入：

import numpy as np

然后，在generate函数之前添加温度采样函数：

def temperature_sampling(logits, temperature=1.0):
    """
    温度采样实现
    :param logits: 模型输出的原始logits
    :param temperature: 温度参数，控制随机性程度
    :return: 采样得到的token id
    """
    # 应用温度缩放
    scaled_logits = logits / temperature
    
    # 转换为概率分布
    probs = softmax(scaled_logits)
    
    # 从概率分布中采样
    next_id = np.random.choice(len(probs), p=probs)
    
    return next_id

修改generate函数支持温度采样

接下来，我们需要修改generate函数以支持温度采样。首先更新函数签名：

def generate(inputs, params, n_head, n_tokens_to_generate, temperature=1.0):
    from tqdm import tqdm
    
    for _ in tqdm(range(n_tokens_to_generate), "generating"):
        logits = gpt2(inputs, **params, n_head=n_head)
        
        # 使用温度采样替代贪婪采样
        next_id = temperature_sampling(logits[-1], temperature)
        
        inputs.append(int(next_id))
    
    return inputs[len(inputs) - n_tokens_to_generate :]

更新main函数参数

最后，更新main函数以接收温度参数：

def main(prompt: str, n_tokens_to_generate: int = 40, model_size: str = "124M", 
         models_dir: str = "models", temperature: float = 1.0):
    from utils import load_encoder_hparams_and_params
    
    encoder, hparams, params = load_encoder_hparams_and_params(model_size, models_dir)
    input_ids = encoder.encode(prompt)
    assert len(input_ids) + n_tokens_to_generate < hparams["n_ctx"]
    
    # 传递温度参数给generate函数
    output_ids = generate(input_ids, params, hparams["n_head"], 
                         n_tokens_to_generate, temperature)
    
    output_text = encoder.decode(output_ids)
    return output_text

添加top-p采样功能

top-p采样（也称为nucleus采样）是另一种先进的采样策略，它从累积概率超过阈值p的最可能token集合中采样。

实现top-p采样函数

在gpt2.py中添加top-p采样函数：

def top_p_sampling(logits, top_p=0.9):
    """
    top-p（nucleus）采样实现
    :param logits: 模型输出的原始logits
    :param top_p: 累积概率阈值
    :return: 采样得到的token id
    """
    # 转换为概率分布
    probs = softmax(logits)
    
    # 按概率降序排序
    sorted_indices = np.argsort(probs)[::-1]
    sorted_probs = probs[sorted_indices]
    
    # 计算累积概率
    cumulative_probs = np.cumsum(sorted_probs)
    
    # 找到累积概率超过top_p的最小token集合
    indices_to_keep = cumulative_probs <= top_p
    if not np.any(indices_to_keep):
        indices_to_keep[0] = True
    
    # 限制token集合
    sorted_indices = sorted_indices[indices_to_keep]
    sorted_probs = sorted_probs[indices_to_keep]
    
    # 重新归一化概率
    sorted_probs = sorted_probs / np.sum(sorted_probs)
    
    # 从限制后的分布中采样
    next_id = np.random.choice(sorted_indices, p=sorted_probs)
    
    return next_id

创建组合采样函数

为了提供最大的灵活性，我们可以创建一个支持多种采样策略的函数：

def sample_from_logits(logits, strategy="greedy", temperature=1.0, top_p=0.9):
    """
    统一的采样函数，支持多种采样策略
    :param logits: 模型输出的原始logits
    :param strategy: 采样策略，可选 "greedy", "temperature", "top_p", "temperature_top_p"
    :param temperature: 温度参数（仅用于温度采样）
    :param top_p: top-p参数（仅用于top-p采样）
    :return: 采样得到的token id
    """
    if strategy == "greedy":
        return np.argmax(logits)
    
    elif strategy == "temperature":
        return temperature_sampling(logits, temperature)
    
    elif strategy == "top_p":
        return top_p_sampling(logits, top_p)
    
    elif strategy == "temperature_top_p":
        # 先应用温度缩放
        scaled_logits = logits / temperature
        # 再应用top-p采样
        return top_p_sampling(scaled_logits, top_p)
    
    else:
        raise ValueError(f"Unknown sampling strategy: {strategy}")

更新generate函数支持多种策略

修改generate函数以支持所有采样策略：

def generate(inputs, params, n_head, n_tokens_to_generate, 
             strategy="greedy", temperature=1.0, top_p=0.9):
    from tqdm import tqdm
    
    for _ in tqdm(range(n_tokens_to_generate), "generating"):
        logits = gpt2(inputs, **params, n_head=n_head)
        
        # 使用指定的采样策略
        next_id = sample_from_logits(logits[-1], strategy, temperature, top_p)
        
        inputs.append(int(next_id))
    
    return inputs[len(inputs) - n_tokens_to_generate :]

完整的使用示例

现在，让我们看看如何使用扩展后的picoGPT：

# 使用默认的贪婪采样
python gpt2.py "人工智能的未来是"

# 使用温度采样，温度=0.8
python gpt2.py "人工智能的未来是" --temperature 0.8

# 使用top-p采样，top_p=0.9
python gpt2.py "人工智能的未来是" --strategy top_p --top_p 0.9

# 使用温度+top-p组合采样
python gpt2.py "人工智能的未来是" --strategy temperature_top_p --temperature 0.7 --top_p 0.95

参数调优建议

不同的采样策略和参数组合会产生不同的效果：

温度参数调优

• temperature=0.1-0.5：保守生成，输出更确定
• temperature=0.5-1.0：平衡的创造性
• temperature=1.0-2.0：高度创造性，可能产生不连贯内容

top-p参数调优

• top_p=0.1-0.5：非常保守，只考虑最可能的token
• top_p=0.7-0.9：推荐范围，平衡质量和多样性
• top_p=0.95-1.0：非常开放，考虑几乎所有token

性能优化技巧

虽然picoGPT本身追求简洁，但我们仍然可以做一些优化：

1. 缓存softmax结果：如果多次调用采样函数，可以缓存softmax计算结果
2. 向量化操作：使用NumPy的向量化操作提高效率
3. 预计算：对于固定的温度值，可以预计算缩放因子

测试你的实现

创建测试脚本验证扩展功能：

# test_sampling.py
import numpy as np

# 测试温度采样
def test_temperature_sampling():
    logits = np.array([1.0, 2.0, 3.0, 4.0])
    
    # 测试不同温度值
    for temp in [0.1, 0.5, 1.0, 2.0]:
        samples = []
        for _ in range(1000):
            sample = temperature_sampling(logits, temp)
            samples.append(sample)
        
        # 验证采样分布
        unique, counts = np.unique(samples, return_counts=True)
        print(f"Temperature={temp}: {dict(zip(unique, counts))}")

# 运行测试
if __name__ == "__main__":
    test_temperature_sampling()

总结

通过本教程，你已经成功为picoGPT添加了温度采样和top-p策略功能。这些扩展使得picoGPT从一个简单的贪婪采样实现变成了一个功能更完整的文本生成工具。记住，采样策略的选择对生成文本的质量和多样性有重要影响，建议根据具体应用场景调整参数。

扩展后的picoGPT现在支持：

• ✅ 贪婪采样（原始功能）
• ✅ 温度采样（控制随机性）
• ✅ top-p采样（控制token集合大小）
• ✅ 温度+top-p组合采样（最佳实践）

这些改进让picoGPT在保持极简代码风格的同时，提供了与大型语言模型相媲美的采样灵活性。现在你可以使用这个增强版的picoGPT来生成更具创造性和多样性的文本内容了！

【免费下载链接】picoGPT An unnecessarily tiny implementation of GPT-2 in NumPy. 项目地址: https://gitcode.com/gh_mirrors/pi/picoGPT