在 TorchScript 中对列表求和的推荐方法是什么

Question

因为 Python builtin <built-in function sum> is currently not supported in Torchscript: 我正在寻找推荐的方法来执行以下操作：

class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return sum(x.tolist())

model = Model()   
model = torch.jit.script(model)     
model(torch.arange(10))

Answer 1

最简单的方法是直接使用 PyTorch 的 sum：

class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return torch.sum(x)

如果出于某种原因，这不是一个选项，您必须使用带有显式循环的类型规范（请注意类型提示！）：

import typing

import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x) -> int:
        x: typing.List[int] = x.tolist()
        result = 0
        for elem in x:
            result += elem
        return result

在 TorchScript 中对列表求和的推荐方法是什么

What is the recommended way to sum a list in TorchScript

jit

pytorch