Function pytorch2libtorch

segmentation/tools/pytorch2torchscript.py:89–129

Export Pytorch model to TorchScript model and verify the outputs are same between Pytorch and TorchScript. Args: model (nn.Module): Pytorch model we want to export. input_shape (tuple): Use this input shape to construct the corresponding dummy input and execute t

(model,
                     input_shape,
                     show=False,
                     output_file='tmp.pt',
                     verify=False)

def pytorch2libtorch(model,
                     input_shape,
                     show=False,
                     output_file='tmp.pt',
                     verify=False):
    """Export Pytorch model to TorchScript model and verify the outputs are
    same between Pytorch and TorchScript.

    Args:
        model (nn.Module): Pytorch model we want to export.
        input_shape (tuple): Use this input shape to construct
            the corresponding dummy input and execute the model.
        show (bool): Whether print the computation graph. Default: False.
        output_file (string): The path to where we store the
            output TorchScript model. Default: `tmp.pt`.
        verify (bool): Whether compare the outputs between
            Pytorch and TorchScript. Default: False.
    """
    if isinstance(model.decode_head, nn.ModuleList):
        num_classes = model.decode_head[-1].num_classes
    else:
        num_classes = model.decode_head.num_classes

    mm_inputs = _demo_mm_inputs(input_shape, num_classes)

    imgs = mm_inputs.pop('imgs')

    # replace the orginal forword with forward_dummy
    model.forward = model.forward_dummy
    model.eval()
    traced_model = torch.jit.trace(
        model,
        example_inputs=imgs,
        check_trace=verify,
    )

    if show:
        print(traced_model.graph)

    traced_model.save(output_file)
    print('Successfully exported TorchScript model: {}'.format(output_file))


def parse_args():