NVIDIA-AI-IOT / torch2trt

An easy to use PyTorch to TensorRT converter

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

[Question & FeatureRequest] Is that possible to convert a torch.nn.Module to TRTMoudle?

ElinLiu0 opened this issue · comments

commented

For example,i got a module class which based on torch.nn.Moudle() like this:

model = VALLE(
        N_DIM,
        NUM_HEAD,
        NUM_LAYERS,
        norm_first=True,
        add_prenet=False,
        prefix_mode=PREFIX_MODE,
        share_embedding=True,
        nar_scale_factor=1.0,
        prepend_bos=True,
        num_quantizers=NUM_QUANTIZERS,
).to(device)

Is that possible convert this module class into a TRTMoudle individually,and still can be using the methods and attributes from VALLE?
I am seen that TRTMoudle has a function load_state_dict(),will it have a simple convert like somehow,uh..
TRTModule.from_torch() or something?

commented

And then just seperate the module initailize and state_dict() loading.