PyTorch文本识别函数返回preds时报UnboundLocalError如何解决

错误原因 #

该报错和函数外部变量赋值无关，触发逻辑如下：

• preds、preds_str的赋值逻辑完全位于for image_tensors, image_path_list in demo_loader循环内部
• 当demo_loader_for_recog()返回的DataLoader没有加载到任何有效样本时，for循环的代码块不会执行，两个变量从未被定义
• 函数执行到return语句时直接引用未定义的变量，就会抛出UnboundLocalError: local variable 'preds' referenced before assignment

控制台输出的DataLoader worker数量警告仅为性能提示：当前Colab环境建议最多开2个worker进程，开4个可能导致数据加载卡顿甚至冻结，和本次变量报错无关联。

修复方法 #

1. 根因修复：检查数据集加载逻辑 #

先确认DataLoader为空的原因，在demo_loader = demo_loader_for_recog()后加一行调试代码：

print(f"加载到的测试样本数量：{len(demo_loader.dataset)}")

如果输出值为0，说明测试图片路径配置错误、目录下无有效图片文件，修正路径配置即可让循环正常执行。

2. 代码容错优化 #

原代码除了变量未定义问题，还存在多batch预测时只返回最后一个batch结果的逻辑bug，优化后代码如下：

def text_recoginition(model_path):
  converter,recognition_model = text_recog_model_config(model_path)
  demo_loader = demo_loader_for_recog()
  recognition_model.eval()
  # 提前初始化结果存储变量，兜底空DataLoader场景
  all_preds = []
  all_preds_str = []
  with torch.no_grad():
      for image_tensors, image_path_list in demo_loader:
          batch_size = image_tensors.size(0)
          image = image_tensors.to(device)
          # 适配最大长度预测的参数初始化
          length_for_pred = torch.IntTensor([opt.batch_max_length] * batch_size).to(device)
          text_for_pred = torch.LongTensor(batch_size, opt.batch_max_length + 1).fill_(0).to(device)
          if 'CTC' in opt.Prediction:
              preds = recognition_model(image, text_for_pred)
              # 贪心解码：取最大概率索引转字符
              preds_size = torch.IntTensor([preds.size(1)] * batch_size)
              _, preds_index = preds.max(2)
              preds_str = converter.decode(preds_index, preds_size)
          else:
              preds = recognition_model(image, text_for_pred, is_train=False)
              # 贪心解码：取最大概率索引转字符
              _, preds_index = preds.max(2)
              preds_str = converter.decode(preds_index, length_for_pred)
          # 逐batch收集结果，避免多batch时丢失前面样本的预测值
          all_preds.append(preds)
          all_preds_str.extend(preds_str)
  # 空数据集场景直接返回空结果，不触发变量未定义错误
  if len(all_preds) == 0:
      return None, []
  # 拼接所有batch的预测张量为完整结果
  final_preds = torch.cat(all_preds, dim=0)
  return final_preds, all_preds_str

内容的提问来源于stack exchange，提问作者Viet Nguyen Hoang