Torchscript/C++ jit::trace 模型 - 访问层参数

Question

我有一个在 python 中训练的模型，使用 torch.jit.trace 进行跟踪，并使用 torch::jit::load.

加载到 C++ 中

有没有办法访问最后一层来提取模型所需输出深度的值（例如，如果它是从 16 -> 2 开始的 Conv2D 层，我想为深度预定义一个张量[b,d->2,x,y] of 2)?

Answer 1

这不是解决这个问题的最优雅的方法，但最直接的方法就是传递一个虚拟张量并访问形状。我尝试过的另一种方法是访问参数列表并寻找“softmax”，不幸的是我不能保证每个人的模型在搜索时都会以相同的方式拼写它。如果其他人对此有很好的答案，请随时分享，但现在必须这样做。