前言

虽然网上有许多大模型可以使用,但输出的结果大多千篇一律,因为可以用来训练这些模型的数据集就那么几个。即便有一些独特的大模型,也多是基于其他人的数据集训练得到的,输出的风格和自己的要求之间还有不小的差距。从零开始训练一个大模型,需要海量的数据,相当长的时间,还有很高的算力。为了得到自己想要的风格,我开始探索如何用微调来快速的获取结果。

《瑞克和莫蒂》是我最喜欢的一部科幻情景喜剧,我喜欢这部作品的原因之一,就是每个角色都有鲜明的性格,而这种性格往往体现在角色的台词上,于是我就有了搭建一个聊天机器人,复刻角色的说话特点的想法,而且还想生成一个瑞克的图片。下文将围绕这两个任务进行,来展示如何让大模型输出我想要的风格。

生成特定风格的文本

如果是ChatGPT的付费用户,OpenAI提供了自制ChatGPT的功能,可以根据给定的文本来生成内容。我至今一直在使用免费的ChatGPT,所以需要另辟蹊径。我使用的是Hugging Face提供的GPT2模型,已经加载了预训练好的权重,只需要一些文本进行微调训练,下面以搭建一个聊天机器人来举例说明。

训练前还需要创建自己的数据集。Hugging Face上有前五季的台词,我做了一下处理,把角色的名字加在了每句台词之前,以便让模型识别出瑞克的说话风格。训练脚本的核心部分只有60行,先加载网上预训练好的模型和分词器,然后对用于微调的数据集进行分词处理,之后就是训练。我使用的设备是NUC 11,它搭载了Nvidia GeForce RTX 2060独立GPU,i7-1165G7处理器,RAM是2×DDR4-3200,使用的OS是Ubuntu 22.04,并且安装了12.1版本的CUDA。在NUC 11上,训练300个循环需要大概半小时的时间。

def train(): 
    # Load GPT-2 model and tokenizer
    tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    config = GPT2Config.from_pretrained("gpt2")
    model = GPT2LMHeadModel.from_pretrained("gpt2", config=config)

    # Create dataset and data_collator
    results_root = os.path.join(os.getcwd(), "data")
    os.makedirs(results_root, exist_ok=True)

    results_path = "/home/nuc/workspace/build-gpt2/data/train.txt"
    train_dataset = TextDataset(
        tokenizer=tokenizer,
        file_path=results_path,
        block_size=128,
    )

    results_path = "/home/nuc/workspace/build-gpt2/data/val.txt"
    val_dataset = TextDataset(
        tokenizer=tokenizer,
        file_path=results_path,
        block_size=128,
    )

    data_collator = create_data_collator(tokenizer)

    # Configure Trainer instance
    training_args = TrainingArguments(
        output_dir="./output",
        overwrite_output_dir=True,
        num_train_epochs=100,
        per_device_train_batch_size=4,
        save_steps=10_000,
        save_total_limit=2,
        logging_steps=100,
    )

    trainer = Trainer(
        model=model,
        args=training_args,
        data_collator=data_collator,
        train_dataset=train_dataset,
        eval_dataset=val_dataset,
    )

    # Train the model
    train_result = trainer.train() 
    metrics = train_result.metrics
    trainer.save_model("./models")
    trainer.log_metrics("train", metrics)
    trainer.save_metrics("train", metrics)

训练部分的代码如上所示,这是完整的代码库。如果对模型的实现细节感兴趣,还可以参考我在B站发布的一系列从零实现LLM的视频。

GPT2只可以做文本的生成,训练集是动画的台词,模型也就只能输出台词,真正的聊天需要使用DialoGPT之类的模型。我不想把这个项目搞得很复杂,这里使用了一个技巧让输出看起来像是在和瑞克对话。动画中,和瑞克互动最多的角色是莫蒂。所以每一个用户的输入,我都在前面加上"Morty: Hi Rick.",造成一种这是莫蒂对瑞克说出的台词的假象。输出的时候,再把"Rick:"去掉,把瑞克对莫蒂的回答,变成对用户说的话。前端使用Streamlit搭建,添加了用户输入文本的输入框,并且以对话的形式展示瑞克的回答,这里是项目页面,可以体验一下聊天机器人的效果。

我问了模型一些简单的问题,比如"你是谁"之类的,模型有时确实会回答说"Rick",而且输出的文本里有结巴的现象,动画中瑞克经常会喝醉,然后变得语无伦次。最能体现瑞克风格的例子无疑是下图。"Wubba Lubba dub-dub"是瑞克的口头禅之一,当他开心或者想要开玩笑的时候就会这么说。

image
模型说出了瑞克的口头禅"Wubba Lubba dub-dub"

生成特定风格的图片

相比文本,平时遇到的更多的需求是如何生成特定风格的图片。处理图片需要更多的算力,NUC 11显得捉襟见肘,所以我使用的是Replicate网站提供的免费算力,试用期只能微调一个模型,之后需要按时间收费。接下来还是收集训练集,我从网上下载了15张动画里瑞克的截图,然后把图片压缩成一个.zip压缩包。Replicate提供了LoRA模型,在下图中灰色方框内上传压缩包即可,相比普通的Stable Diffusion模型,LoRA训练时间更短,模型也更小。模型的训练任务有三种可选的类型,分别是人脸,物体,和风格,这里选择的是风格,其他参数不需要改动。

在灰色的方框内,上传自己的训练集

几分钟后训练完毕,在页面的右侧会看到类似下图的内容。右键点击Preview下方带有下载图标的方框,然后复制链接。打开一个可以接收微调权重的模型,比如这个,这一步会提示要登录才能跑模型,有GitHub账户的话可以直接登录。

接着在prompt一栏中输入类似a photo of an astronaut riding a horse in the style of <1>的语句,其中<1>表示的是在下面lora_urls一栏里输入的模型,也就是刚才训练好模型后,权重的链接。

在prompt(提示词)一栏中输入提示词, nagative_prompt是不想让模型生成什么的屏蔽词汇,同时还需要把微调后得到的权重输入到lora_urls一栏种
output
模型根据我的提示词,输出的一个穿着绿色衣服的瑞克。我使用的训练集比较小,模型的可玩度不是很高,所以输出也比较简单,这里的结果仅仅是一个示范

也可以训练多个模型,每个模型在一个特定风格上微调,把输出的模型放在一起使用,这时<1>, <2>, <3>...分别表示不同模型的代号, 还要在lora_urls一栏用竖杠符号分割链接。如果觉得效果不够理想,可以调整lora_scales一栏的值,默认是0.5,可以增大或者减小试试。

结语

本文介绍了如何通过大模型的微调快速获得特定风格,并列举了文本生成和图片生成的两个实例作为参考。随着flash attention和英伟达Apex工具的出现,模型的训练越来越快,但对算力的要求还是很高。我尝试过类似Sora的大模型进行视频生成,但需要好几张A100显卡才能跑通,内存要求在24GB以上,我的显卡内存只有8GB,远远达不到要求,所以还没找到合适的方法微调生成视频的模型。另外,本文介绍的方式都是免费的,付费的方式也留待之后再去探索。