PyTorch转ScriptModule的问题记录

本文记录了转ScriptModule时遇到的一系列问题

如何转ScriptModule?

这里有提到背景

https://blog.csdn.net/ReadyShowShow/article/details/137275857

遇到的坑

Expected a value of type ‘Tensor’ for argument ‘self’ but instead found type ‘Optional[Tensor]’.

数据类型为可None类型，但算子不接受None，需要去掉类型的可None。

xs:Optional[torch.Tensor] = None
assert xs is not None
x = xs / 10

Expected integer literal for index but got a variable or non-integer. ModuleList/Sequential indexing is only supported with integer literals. For example, ‘i = 4; self.layersi’ will fail because i is not a literal. Enumeration is supported, e.g. ‘for index, v in enumerate(self): out = v(inp)’:

Attribute 1 is not of annotated type

通常是由 ModuleList 或者 ModuleDict 的遍历或取元素引起的。

3 种解法：
第一种是推荐做法: 直接使用遍历

m = nn.ModuleList()
m.append()
for index, v in enumerate(m):
  out = v(inp)

第二种是无法直接遍历，必须用下标取值。但是自定义的module，可以把继承torch.nn.Module改为继承torch.jit.ScriptModule，下标取值后的类型明确需要，且显示调用forward

class  MultiHeadAttention(torch.jit.ScriptModule):
	def __init__(self):
	  ...
	@torch.jit.script_method
    def forward(self, ):
      ...

m = nn.ModuleList()
m.append(MultiHeadAttention())

for i in range(m):
  l:MultiHeadAttention = m[i] # 明确类型
  r = l.forward(xxx)

第三种也是无法直接遍历，必须用下标