问题描述

模型里面有lambda函数，在保存模型的时候出现了如下问题：

Traceback (most recent call last):
  File "D:\MaxViT_Demo\train.py", line 226, in <module>
    val_list, pred_list, val_loss, val_acc = val(model_ft, DEVICE, test_loader)
  File "D:\Users\wh109\anaconda3\lib\site-packages\torch\autograd\grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "D:\MaxViT_Demo\train.py", line 115, in val
    torch.save(model, file_dir + "/" + 'model_' + str(epoch) + '_' + str(round(acc, 3)) + '.pth')
  File "D:\Users\wh109\anaconda3\lib\site-packages\torch\serialization.py", line 380, in save
    _save(obj, opened_zipfile, pickle_module, pickle_protocol)
  File "D:\Users\wh109\anaconda3\lib\site-packages\torch\serialization.py", line 590, in _save
    pickler.dump(obj)
AttributeError: Can't pickle local object 'Stage.__init__.<locals>.<lambda>'

Process finished with exit code 1

解决办法

打开serialization.py文件
将

import pickle

替换为：

import dill as pickle