Working with DLRM in onnxruntime.

[atmadeep]

Hi, I'm working on getting the DLRM 10GB model to work with onnxruntime. I had the following questions:

1. I've exported the model from pytorch scripts using various ways.
   a. Added a torch.onnx.export in the dlrm_pytorch script to export the model as soon as it's initialized. The model size was 85kb.
   b. Instantiated a model in pytorch ( in a standalone script) and loaded the weights. Used torch.onnx.export. Failed due to 2GB export limit in torch.

import sys
import numpy as np
import torch

sys.path.insert(0, '/home/user/MLPerfBench/DLRMBench/Scripts/python/')

from dlrm_model import DLRM_Net
torch.__version__
device = torch.device("cpu")
m_spa=64
ln_emb=np.array([9980333,36084,17217,7378,20134,3,7112,1442,61,9758201,1333352,313829,10,2208,11156,122,4,970,14,9994222,7267859,9946608,415421,12420,101,36])
ln_bot=np.array([13,512,256,64])
ln_top=np.array([415,512,512,256,1])

dlrm = DLRM_Net(
            m_spa,
            ln_emb,
            ln_bot,
            ln_top,
            arch_interaction_op="dot",
            arch_interaction_itself=False,
            sigmoid_bot=-1,
            sigmoid_top=ln_top.size - 2,
            sync_dense_params=True,
            loss_threshold=0.0,
            ndevices=-1,
            qr_flag=False,
            qr_operation=None,
            qr_collisions=None,
            qr_threshold=None,
            md_flag=False,
            md_threshold=None,
        )
model_path = "/home/ext/atmaarya/my_work/MLPerfBench/DLRMBench/Models/dlrm_terabyte.pytorch"
ld_model = torch.load(model_path)

dlrm.load_state_dict(ld_model["state_dict"])
dlrm = dlrm.to(device)
dlrm.eval()
torch.save(dlrm, "dlrm_model.pt")
dlrm_model_inputs = torch.load("/home/ext/atmaarya/my_work/dlrm_model_inputs")
x = dlrm_model_inputs['x']
IS_o = dlrm_model_inputs['lS_o']
IS_i = dlrm_model_inputs['lS_i']
symbolic_names = {0:'batch_size'}
x_new = x[:1,:]
IS_o_new = IS_o[:,:1]
IS_i_new = IS_i[:,:1]
print(x_new)
print(IS_o[])

torch.onnx.export(dlrm, args=(x_new,IS_o_new, IS_i_new), f="dlrm_model.onnx", opset_version=11, input_names= ['x','IS_o','IS_i'], output_names=['pred'],use_external_data_format=True)

c. Downloaded the onnx.tar file from [here] (https://dlrm.s3-us-west-1.amazonaws.com/models/tb0875_10M.onnx.tar). I've got the top, bottom and embedding weights. But the model obtained is still 37kb.

So is the model size in this supposed to be this small?

2. Is there a working example of this model in onnxruntime? I've worked previously with other model, but with DLRM, how am I supposed to use the weight files is the confusing part. Any help would be appreciated?