Discord Bot加载Torch模型触发AttributeError问题求助

我想在Discord机器人中加载PyTorch训练好的GAN生成器（Generator）和判别器（Discriminator）模型。直接运行bot.py时模型加载正常，但执行python main.py时出现错误：

AttributeError: Can't get attribute 'Generator' on <module '__main__' from '/home/mle/Devel/gan-discord-bot/main.py'>

错误原因 #

PyTorch的torch.load()加载模型时，会默认从**当前主模块（__main__）**查找模型类的定义：

• 直接运行bot.py时，bot.py就是__main__模块，它已经导入了model.py中的Generator和Discriminator，因此能正常加载；
• 运行main.py时，main.py成为__main__模块，它仅导入了bot模块，torch.load()会在main.py中寻找Generator类，自然找不到，触发报错。

解决方案 #

方法1：分离模型结构与权重加载（推荐） #

这种方式是PyTorch的最佳实践，仅加载模型权重而非整个模型对象，彻底避开模块命名空间冲突：

修改bot.py中的模型加载代码：

import discord
import torch
from torchvision.utils import save_image

from model import *

# 先初始化模型结构（ngpu参数要和训练时一致，训练用GPU填对应数量，CPU填0）
netG = Generator(ngpu=0)
# 加载预训练权重
netG.load_state_dict(torch.load(netG_name, map_location=torch.device('cpu')))

netD = Discriminator(ngpu=0)
netD.load_state_dict(torch.load(netD_name, map_location=torch.device('cpu')))

netG.eval()
netD.eval()

# 后续generate_image、run_discord_bot函数保持不变

方法2：让主模块导入模型类 #

在main.py中直接导入model模块的内容，让__main__模块包含模型类定义：

import bot
from model import *  # 导入模型类，使__main__模块拥有相关定义

if __name__=='__main__':
    bot.run_discord_bot()

缺点是会让main.py引入不必要的依赖，不够优雅。

方法3：调整执行入口 #

将run_discord_bot()的调用移到bot.py中，直接以bot.py作为执行入口，不再使用main.py：

# 在bot.py末尾添加
if __name__ == '__main__':
    run_discord_bot()

之后运行python bot.py即可，完全避开__main__模块的冲突问题。

推荐选择 #

优先使用方法1，它不仅解决了当前问题，还让代码结构更清晰，后续修改模型结构或迁移代码时更不容易出现兼容问题。

内容的提问来源于stack exchange，提问作者L Maxime