Modifying interaction layer to include 2 MLPs in DLRM

facebookresearch / recipes

Recipes are a standard, well supported set of blueprints for machine learning engineers to rapidly train models using the latest research techniques without significant engineering overhead.Specifically, recipes aims to provide- Consistent access to pre-trained SOTA models ready for production- Reference implementations for SOTA research reproducibility, and infrastructure to guarantee correctness, efficiency, and interoperability.

BSD 3-Clause "New" or "Revised" License

284 stars 35 forks source link

Summary: X-link: https://github.com/pytorch/torchrec/pull/382

X-link: https://github.com/facebookresearch/dlrm/pull/242

This diff adds 2 MLPs to the interaction layer in DLRM for MLPerf update. New DLRM module called DLRMV2 can be realized by --dlrmv2 argument. Additional arguments for the interaction MLPs are --interaction_branch1_layer_sizes and --interaction_branch2_layer_sizes to pass in the MLP sizes. The output dimension of the interaction MLPs must be a multiple of the embedding dimension.

DLRMTrain now takes in a DLRM/DLRMV2 module at construction time.

Reviewed By: colin2328, samiwilf

Differential Revision: D35861688

facebookresearch / recipes

Modifying interaction layer to include 2 MLPs in DLRM #26