def get_dlrm_model(sparse_dlrm=False):
"""Obtain dlrm model. The configs specified are based on the script in
bench/dlrm_s_criteo_kaggle.sh. The same config is used to train the model
for benchmarking on data sparsifier.
"""
dlrm_model_config = {
'm_spa':
16,
'ln_emb':
np.array([
1460, 583, 10131227, 2202608, 305, 24, 12517, 633, 3, 93145, 5683,
8351593, 3194, 27, 14992, 5461306, 10, 5652, 2173, 4, 7046547, 18,
15, 286181, 105, 142572
],
dtype=np.int32),
'ln_bot':
np.array([13, 512, 256, 64, 16]),
'ln_top':
np.array([367, 512, 256, 1]),
'arch_interaction_op':
'dot',
'arch_interaction_itself':
False,
'sigmoid_bot':
-1,
'sigmoid_top':
2,
'sync_dense_params':
True,
'loss_threshold':
0.0,
'ndevices':
1,
'qr_flag':
False,
'qr_operation':
'mult',
'qr_collisions':
4,
'qr_threshold':
200,
'md_flag':
False,
'md_threshold':
200,
'weighted_pooling':
None,
'loss_function':
'bce'
}
if sparse_dlrm:
dlrm_model = SparseDLRM(**dlrm_model_config)
else:
dlrm_model = DLRM_Net(**dlrm_model_config)
return dlrm_model
ln_top = np.array([479, 1024, 1024, 512, 256, 1])
else:
raise ValueError("only --max-in-range 10M or 40M is supported")
else:
raise ValueError("only kaggle|terabyte dataset options are supported")
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,
)
# Load model is specified
if not (args.load_model == ""):
print("Loading saved model {}".format(args.load_model))
ld_model = torch.load(args.load_model,
def load(self, model_path, inputs=None, outputs=None):
# debug prints
# print(model_path, inputs, outputs)
dlrm = DLRM_Net(
self.m_spa,
self.ln_emb,
self.ln_bot,
self.ln_top,
arch_interaction_op="dot",
arch_interaction_itself=False,
sigmoid_bot=-1,
sigmoid_top=self.ln_top.size - 2,
sync_dense_params=True,
loss_threshold=0.0,
ndevices=self.ndevices,
qr_flag=False,
qr_operation=None,
qr_collisions=None,
qr_threshold=None,
md_flag=False,
md_threshold=None,
)
if self.use_gpu:
dlrm = dlrm.to(self.device) # .cuda()
if dlrm.ndevices > 1:
dlrm.emb_l = dlrm.create_emb(self.m_spa, self.ln_emb)
if self.use_gpu:
if dlrm.ndevices > 1:
# NOTE: when targeting inference on multiple GPUs,
# load the model as is on CPU or GPU, with the move
# to multiple GPUs to be done in parallel_forward
ld_model = torch.load(model_path)
else:
# NOTE: when targeting inference on single GPU,
# note that the call to .to(device) has already happened
ld_model = torch.load(
model_path,
map_location=torch.device('cuda')
# map_location=lambda storage, loc: storage.cuda(0)
)
else:
# when targeting inference on CPU
ld_model = torch.load(model_path, map_location=torch.device('cpu'))
# debug print
# print(ld_model)
dlrm.load_state_dict(ld_model["state_dict"])
self.model = dlrm
# find inputs from the model if not passed in by config
if inputs:
self.inputs = inputs
else:
self.inputs = []
initializers = set()
for i in self.model.graph.initializer:
initializers.add(i.name)
for i in self.model.graph.input:
if i.name not in initializers:
self.inputs.append(i.name)
# find outputs from the model if not passed in by config
if outputs:
self.outputs = outputs
else:
self.outputs = []
for i in self.model.graph.output:
self.outputs.append(i.name)
# debug print
# for e in self.ln_emb:
# print('Embedding', type(e), e)
return self