keyerror state_dict

在打通信比赛的时候，想尝试将fedformer框架应用到无线信道反馈场景上来，因此需要对官方的baseline做很大的修改。在模型存储的时候遇到了一些问题在此记录以下。
官方中有一段代码：

model_encoder1.load_state_dict(torch.load(path)['state_dict'])

这种写法很奇怪，之前从来没见过。之前一般都是这样写的：

model_encoder1.load_state_dict(torch.Load(path).state_dict())

第一次见到加载模型的时候可以使用下标运算符访问。所以尝试将训练好的模型用官方baseline加载却出现了如下错误：object is not subscriptable。
当时我的存储模型代码为：torch. Save(model, PATH)。当时以为是因为存储的不是字典类型所以没有下标运算符。所以修改存储代码为：torch.save(model.state_dict(), PATH)
直接存储为字典类型，结果出现了这个错误：keyerror state_dict。这是意思是没有state_dict这个关键字。在网上查询了许多资料都没有很好的解决。最终在pytroch官网上发现了解决办法。原来pytorch在存储的时候是可以把loss、epoch这些设置与模型一起存储为一个字典类型。以下为官方示例代码：

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

之后加载模型时可以使用类似字典访问的办法来进行访问：

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model. Train()

这样对复现别人论文十分有帮助，甚至pytorch可以把多个模型存储在一个文件中：

torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)

加载模型：

modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

可以看到pytorch还是十分强大的，以后查资料就直接去访问官网就好了。