nonebot2聊天机器人插件8：基于GAN的虚拟巡天quasistar_engine
该插件涉及知识点：图片文件读写冲突，对抗式生成网络，GPU模型转CPU运算
插件合集：nonebot2聊天机器人插件
该系列为用于QQ群聊天机器人的nonebot2相关插件，不保证完全符合标准规范写法，如有差错和改进余地，欢迎大佬指点修正。
前端：nonebot2
后端：go-cqhttp
插件所用语言：python3
前置环境安装过程建议参考零基础2分钟教你搭建QQ机器人——基于nonebot2，但是请注意该教程中的后端版本过旧导致私聊发图异常，需要手动更新go-cqhttp版本。
1. 插件用途
插件名：类星引擎（Quasistar Engine），或者虚拟巡天~
将训练好的多个对抗式生成网络模型安装到插件内，随机从多个模型中选择一个，通过随机噪声生成星球图片。
1、每个帐号每天都可以得到一张固定的星球图片，随机生成结果每日轮换。
2、bot管理员能够用指定字符串作为种子，生成指定星球。
由于图片由深度学习的对抗式生成网络绘制，因此除非哈希值相同，否则几乎不可能出现完全一样的图片（但受限于训练集数量，可能看起来很相似）。
插件需要能够在没有GPU的云服务器上运行，使用CPU运算。
注：与之前的插件不同，这个插件需要训练好的生成器模型才能工作，训练相关模型文件需要具备神经网络方面的知识，在之前AI相关的博文中给出过该插件中用到的模型文件的训练方法，也不排除之后进一步改进的可能。
2. 目录结构
在plugins文件夹中新建一个文件夹quasistar_engine，文件夹内目录结构如下：
|-quasistar_engine
|-generator_models
|-所有生成器模型的.pth参数文件
|-temp_img
|-临时图片存储位置
|-__init__.py
|-quasistar_engine.py
|-config.py
|-run_generator.py
|-model.py
其中temp_img为用于存储发送的临时图片的文件夹，generator_models为储存生成器模型数据文件的文件夹，quasistar_engine.py为命令事件响应器的位置，config.py用于存储配置项，run_generator.py用于根据参数运行模型生成图片，model.py为pytorch的生成器网络架构，__init__.py为程序启动位置。
3. 实现难点与解决方案
3.1 图片文件读写冲突
如果与之前的插件一样采用固定临时图片名发送，那么在不同群同时发送命令时，可能会产生图片文件的读写冲突导致异常，因此对于不同群组和私聊的信息，使用不同的文件名命名图片文件。
3.2 对抗式生成网络
对抗式生成网络的训练方法在之前的博文里面已经有详细解释：
彩色星球图片生成1：使用Gan实现（pytorch版）
彩色星球图片生成2：同时使用传统Gan判别器和马尔可夫判别器（pytorch版）
彩色星球图片生成3：代码改进（pytorch版）
彩色星球图片生成4：转置卷积+插值缩放+卷积收缩（pytorch版）
此处不再过多赘述，后续仍然可能有进一步改进的博客。
3.3 GPU模型转CPU运算
网络参数模型为了速度，是在GPU环境下训练的，而GPU云服务器虽然快，却价格过于昂贵，没有足够的财力配备所以说挖矿的真是害人 。
因此在读取网络模型参数时，需要将GPU训练的参数文件转换为使用CPU运算的格式。
将参数加载代码从G_model.load_state_dict(load(G_model_path))
改为G_model.load_state_dict(load(G_model_path, map_location=device('cpu')))。
4. 代码实现
__init__.py
from .quasistar_engine import *
config.py
class Config:
# 记录在哪些群组中使用
used_in_group = ["131551175"]
# 插件执行优先级
priority = 10
# 机器人QQ号
bot_id = "123456789"
# 管理员QQ号，管理员无视冷却cd
super_uid = ["673321342"]
# 触发冷却时间（秒），在这段时间内不会连续两次触发
cd = 10
model.py
import torch.nn as nn
# 生成器,此处填入模型文件对应的生成器网络结构
class G_net(nn.Module):
def __init__(self):
super(G_net, self).__init__()
pass
def forward(self, img_seeds):
pass
return imgs
# 获取模型
def get_G_model():
model = G_net()
return model
run_generator.py
import random
from .model import get_G_model
import os
from time import sleep
from torchvision.utils import save_image
import torch
# 存储所有模型的地址
models_path = os.path.split(os.path.realpath(__file__))[0] + '/generator_models'
# 存储所有临时图片的地址
temp_img_path = os.path.split(os.path.realpath(__file__))[0] + '/temp_img/'
# 存储所有待用的生成器模型
G_models = []
# 初始化加载所有生成器模型参数
def init_G_models(G_models):
# 首先，读出所有模型文件的存储目录【允许任意个模型文件】
# 读取文件夹下所有数据文件
for path, dirs, files in os.walk(models_path, topdown=False):
for name in files:
G_models.append(os.path.join(path, name))
# 对每一个模型文件，都创建一个生成器模型，并且保存在内存中
def path_to_model(G_model_path):
G_model = get_G_model()
# 从磁盘加载之前保存的模型参数
try:
G_model.load_state_dict(torch.load(G_model_path, map_location=torch.device('cpu')))
return G_model
except:
print("模型数据文件" + G_model_path + "不存在或使用了错误的文件,请关闭程序并检查文件，一分钟后程序将自动关闭。")
sleep(60)
exit()
# 所有路径字符串都转换为了生成器模型
G_models = list(map(path_to_model, G_models))
# 需要设置为训练模式，以通过dropout产生更多的GAN多样性
# 但是BN层必须被设置为测试模式，以保证输出性能
for G_model in G_models:
G_model.train()
for _, module in G_model.named_modules():
# 将所有BN层设置为测试模式
if isinstance(module, nn.BatchNorm1d) or isinstance(module, nn.BatchNorm2d):
module.training = False
return G_models
# 初始化加载所有生成器模型参数
G_models = init_G_models(G_models)
# 创建图片
# 输入种子数和图片名，将图片写入磁盘，返回一张由GAN生成的图片的地址
def create_img(seed_int, img_name):
# 生成器输入种子数的维度
img_seed_dim = 128
# 图片分辨率
img_size = 528
# 输出图片的路径
img_path = temp_img_path + img_name + '.png'
# 固定随机数种子
random.seed(seed_int)
torch.manual_seed(seed_int)
# 根据种子数随机选择一个生成器模型使用
G_model = random.choice(G_models)
with torch.no_grad():
# 输出一张生成器产生的图片到输出文件夹
# 产生随机正态分布噪声
img_seeds = torch.randn(1, img_seed_dim)
# 生成图像
fake_images = G_model(img_seeds)
# 将图像缩放到[0,1]的区间
fake_images = 0.5 * (fake_images + 1)
fake_images = fake_images.clamp(0, 1)
# 用torchvision自带的save_image()函数输出到磁盘文件
fake_images = fake_images.view(-1, 3, img_size, img_size)
save_image(fake_images, img_path)
# 返回输出图像的路径
return img_path
quasistar_engine.py
from nonebot import on_command
from nonebot.typing import T_State
from nonebot.adapters import Bot, Event
from nonebot.adapters.cqhttp import MessageSegment
from .config import Config
from time import time, localtime
import json
from hashlib import md5
from .run_generator import *
from nonebot.permission import SUPERUSER
# 记录上一次响应时间
last_response = {}
# 判断是否过了响应cd的函数，默认使用配置文件中的cd
# 如果已经超过了最短响应间隔，返回True
def cool_down(group_id, cd = Config.cd):
global last_response
if group_id not in last_response:
return True
else:
return time() - last_response[group_id] > cd
# 每个用户都可以得到的每日随机星球图片
search_star = on_command("观星", priority=Config.priority)
@search_star.handle()
async def handle_first_receive(bot: Bot, event: Event, state: T_State):
ids = event.get_session_id()
allow_use = True
# 如果这是一条群聊信息
if ids.startswith("group"):
_, group_id, user_id = event.get_session_id().split("_")
if group_id not in Config.used_in_group:
allow_use = False
# 对于私聊信息，在前方加上private_作为冷却时间存储的key
else:
user_id = ids
group_id = 'private_'+ids
# 如果允许使用
if allow_use:
# 如果已经过了冷却时间，或者用户是管理员
if cool_down(group_id) or user_id in Config.super_uid:
# 如果用户不是超级用户，更新cd时间
if user_id not in Config.super_uid:
last_response[group_id] = time()
# 对超级用户，直接使用super_qq号作为存储key与图片保存名，防止在cd期间内出现文件读写冲突
if user_id in Config.super_uid:
group_id = 'super_'+user_id
# 获取用户昵称
infos = str(await bot.get_stranger_info(user_id=user_id))
nickname = json.loads(infos.replace("'", '"'))['nickname']
# 昵称与用户名的组合
nickname_and_id = nickname + '(' + str(user_id) + ')'
# 使用[QQ号+当日日期]作为随机数种子
# 获取日期
local = localtime(time())
today = f"{local[0]}年{local[1]}月{local[2]}日"
# 组合种子字符串
seed_str = user_id + today
# 将种子字符串转换为种子数
# text to md5
md5_str = md5(seed_str.encode('utf-8'))
# md5 to int
seed_int = int(str(str(int('0x' + md5_str.hexdigest(), 0)))[-16:])
# 输入种子数和图片名创建图片，使用group_id作为图片名防止冲突
img_path = create_img(seed_int, group_id)
# 将图片发送给用户
await search_star.send(f"{nickname_and_id}\n{today}\n你观测到了星球：\n"+MessageSegment.image('file:///' + img_path))
# 发送完之后删除临时文件夹中的图片
os.remove(img_path)
# bot管理员专属命令，根据后缀信息创造星球
create_star = on_command("创星", permission=SUPERUSER, priority=Config.priority)
@create_star.handle()
async def handle_first_receive(bot: Bot, event: Event, state: T_State):
msg = str(event.get_message()).strip().replace('\r\n', '').replace('\n', '')
ids = event.get_session_id()
# 如果这是一条群聊信息
if ids.startswith("group"):
_, group_id, user_id = event.get_session_id().split("_")
# 对于私聊信息，在前方加上private_作为冷却时间存储的key
else:
user_id = ids
# 将种子字符串转换为种子数
# text to md5
md5_str = md5(msg.encode('utf-8'))
# md5 to int
seed_int = int(str(str(int('0x' + md5_str.hexdigest(), 0)))[-16:])
# 输入种子数和图片名创建图片，使用group_id作为图片名防止冲突
img_path = create_img(seed_int, "create_star")
# 获取用户昵称
infos = str(await bot.get_stranger_info(user_id=user_id))
nickname = json.loads(infos.replace("'", '"'))['nickname']
# 昵称与用户名的组合
nickname_and_id = nickname + '(' + str(user_id) + ')'
# 将图片发送给用户
await search_star.send(f"{nickname_and_id}\n你创造了{msg}星：\n"+MessageSegment.image('file:///' + img_path))
# 发送完之后删除临时文件夹中的图片
os.remove(img_path)
# 查询帮助命令
star_helper = on_command("巡天帮助", priority=Config.priority)
@star_helper.handle()
async def handle_first_receive(bot: Bot, event: Event, state: T_State):
await star_helper.finish('''巡天指令说明
普通用户指令：
观星——每天创建一颗随机星球，当天结果固定
管理员指令：
创星 [星球名字]——根据星球名字创建一颗随机星球''')
5. 插件配图
该插件无配图
6. 实际效果
扫一扫
824
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
