Skip to content

GitLab

Commit c43a53e4 authored by sunxx1's avatar sunxx1
Browse files

Merge branch 'add_Recommendation' into 'main' 

添加VAE-CF和dlrm

See merge request dcutoolkit/deeplearing/dlexamples_new!24
parents 5394b117 56225fdf
Hide whitespace changes
Inline Side-by-side
Showing with 8769 additions and 0 deletions
PyTorch/Recommendation/VAE-CF/vae/utils/round.py 0 → 100644
View file @ c43a53e4
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import numpy as np
from functools import partial
def round_n(x, n=8):
return n * int(np.ceil(x / n))
round_8 = partial(round_n, n=8)
PyTorch/Recommendation/dlrm/.gitignore 0 → 100644
View file @ c43a53e4
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# PEP 582; used by e.g. github.com/David-OConnor/pyflow
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
PyTorch/Recommendation/dlrm/CODE_OF_CONDUCT.md 0 → 100644
View file @ c43a53e4
# Code of Conduct
Facebook has adopted a Code of Conduct that we expect project participants to adhere to.
Please read the [full text](https://code.fb.com/codeofconduct/)
so that you can understand what actions will and will not be tolerated.
PyTorch/Recommendation/dlrm/CONTRIBUTING.md 0 → 100644
View file @ c43a53e4
# Contributing to DLRM
We want to make contributing to this project as easy and transparent as
possible.
## Pull Requests
We actively welcome your pull requests.
1. Fork the repo and create your branch from `main`.
2. If you've added code that should be tested, add tests.
3. If you've changed APIs, update the documentation.
4. Ensure the test suite passes.
5. Make sure your code lints.
6. If you haven't already, complete the Contributor License Agreement ("CLA").
## Contributor License Agreement ("CLA")
In order to accept your pull request, we need you to submit a CLA. You only need
to do this once to work on any of Facebook's open source projects.
Complete your CLA here: <https://code.facebook.com/cla>
## Issues
We use GitHub issues to track public bugs. Please ensure your description is
clear and has sufficient instructions to be able to reproduce the issue.
Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
disclosure of security bugs. In those cases, please go through the process
outlined on that page and do not file a public issue.
## Coding Style
* 4 spaces for indentation rather than tabs
* 80 character line length
* in general, please maintain a consistent style with the rest of the code
## License
By contributing to DLRM, you agree that your contributions will be licensed
under the LICENSE file in the root directory of this source tree.
PyTorch/Recommendation/dlrm/Dockerfile 0 → 100644
View file @ c43a53e4
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
ARG FROM_IMAGE_NAME=pytorch/pytorch:1.3-cuda10.1-cudnn7-runtime
FROM ${FROM_IMAGE_NAME}
ADD requirements.txt .
RUN pip install -r requirements.txt
RUN pip install torch==1.3.1
WORKDIR /code
ADD . .
PyTorch/Recommendation/dlrm/LICENSE 0 → 100644
View file @ c43a53e4
MIT License
Copyright (c) Facebook, Inc. and its affiliates.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
PyTorch/Recommendation/dlrm/README.md 0 → 100644
View file @ c43a53e4
Deep Learning Recommendation Model for Personalization and Recommendation Systems:
=================================================================================
## 模型结构
```
output:
probability of a click
model: |
/\
/__\
|
_____________________> Op <___________________
/ | \
/\ /\ /\
/__\ /__\ ... /__\
| | |
| Op Op
| ____/__\_____ ____/__\____
| |_Emb_|____|__| ... |_Emb_|__|___|
input:
[ dense features ] [sparse indices] , ..., [sparse indices]
```
More precise definition of model layers:
1) fully connected layers of an mlp
z = f(y)
y = Wx + b
2) embedding lookup (for a list of sparse indices p=[p1,...,pk])
z = Op(e1,...,ek)
obtain vectors e1=E[:,p1], ..., ek=E[:,pk]
3) Operator Op can be one of the following
Sum(e1,...,ek) = e1 + ... + ek
Dot(e1,...,ek) = [e1'e1, ..., e1'ek, ..., ek'e1, ..., ek'ek]
Cat(e1,...,ek) = [e1', ..., ek']'
where ' denotes transpose operation
测试用例执行
--------------------
1) 模型简单测试
```
$ python dlrm_s_pytorch.py --mini-batch-size=2 --data-size=6
time/loss/accuracy (if enabled):
Finished training it 1/3 of epoch 0, -1.00 ms/it, loss 0.451893, accuracy 0.000%
Finished training it 2/3 of epoch 0, -1.00 ms/it, loss 0.402002, accuracy 0.000%
Finished training it 3/3 of epoch 0, -1.00 ms/it, loss 0.275460, accuracy 0.000%
```
2) debug模式（可以自行设置模型参数、规格）
```
$ python dlrm_s_pytorch.py --mini-batch-size=2 --data-size=6 --debug-mode
model arch:
mlp top arch 3 layers, with input to output dimensions:
[8 4 2 1]
# of interactions
8
mlp bot arch 2 layers, with input to output dimensions:
[4 3 2]
# of features (sparse and dense)
4
dense feature size
4
sparse feature size
2
# of embeddings (= # of sparse features) 3, with dimensions 2x:
[4 3 2]
data (inputs and targets):
mini-batch: 0
[[0.69647 0.28614 0.22685 0.55131]
[0.71947 0.42311 0.98076 0.68483]]
[[[1], [0, 1]], [[0], [1]], [[1], [0]]]
[[0.55679]
[0.15896]]
mini-batch: 1
[[0.36179 0.22826 0.29371 0.63098]
[0.0921 0.4337 0.43086 0.49369]]
[[[1], [0, 2, 3]], [[1], [1, 2]], [[1], [1]]]
[[0.15307]
[0.69553]]
mini-batch: 2
[[0.60306 0.54507 0.34276 0.30412]
[0.41702 0.6813 0.87546 0.51042]]
[[[2], [0, 1, 2]], [[1], [2]], [[1], [1]]]
[[0.31877]
[0.69197]]
initial parameters (weights and bias):
[[ 0.05438 -0.11105]
[ 0.42513 0.34167]
[-0.1426 -0.45641]
[-0.19523 -0.10181]]
[[ 0.23667 0.57199]
[-0.16638 0.30316]
[ 0.10759 0.22136]]
[[-0.49338 -0.14301]
[-0.36649 -0.22139]]
[[0.51313 0.66662 0.10591 0.13089]
[0.32198 0.66156 0.84651 0.55326]
[0.85445 0.38484 0.31679 0.35426]]
[0.17108 0.82911 0.33867]
[[0.55237 0.57855 0.52153]
[0.00269 0.98835 0.90534]]
[0.20764 0.29249]
[[0.52001 0.90191 0.98363 0.25754 0.56436 0.80697 0.39437 0.73107]
[0.16107 0.6007 0.86586 0.98352 0.07937 0.42835 0.20454 0.45064]
[0.54776 0.09333 0.29686 0.92758 0.569 0.45741 0.75353 0.74186]
[0.04858 0.7087 0.83924 0.16594 0.781 0.28654 0.30647 0.66526]]
[0.11139 0.66487 0.88786 0.69631]
[[0.44033 0.43821 0.7651 0.56564]
[0.0849 0.58267 0.81484 0.33707]]
[0.92758 0.75072]
[[0.57406 0.75164]]
[0.07915]
DLRM_Net(
(emb_l): ModuleList(
(0): EmbeddingBag(4, 2, mode=sum)
(1): EmbeddingBag(3, 2, mode=sum)
(2): EmbeddingBag(2, 2, mode=sum)
)
(bot_l): Sequential(
(0): Linear(in_features=4, out_features=3, bias=True)
(1): ReLU()
(2): Linear(in_features=3, out_features=2, bias=True)
(3): ReLU()
)
(top_l): Sequential(
(0): Linear(in_features=8, out_features=4, bias=True)
(1): ReLU()
(2): Linear(in_features=4, out_features=2, bias=True)
(3): ReLU()
(4): Linear(in_features=2, out_features=1, bias=True)
(5): Sigmoid()
)
)
time/loss/accuracy (if enabled):
Finished training it 1/3 of epoch 0, -1.00 ms/it, loss 0.451893, accuracy 0.000%
Finished training it 2/3 of epoch 0, -1.00 ms/it, loss 0.402002, accuracy 0.000%
Finished training it 3/3 of epoch 0, -1.00 ms/it, loss 0.275460, accuracy 0.000%
updated parameters (weights and bias):
[[ 0.0543 -0.1112 ]
[ 0.42513 0.34167]
[-0.14283 -0.45679]
[-0.19532 -0.10197]]
[[ 0.23667 0.57199]
[-0.1666 0.30285]
[ 0.10751 0.22124]]
[[-0.49338 -0.14301]
[-0.36664 -0.22164]]
[[0.51313 0.66663 0.10591 0.1309 ]
[0.32196 0.66154 0.84649 0.55324]
[0.85444 0.38482 0.31677 0.35425]]
[0.17109 0.82907 0.33863]
[[0.55238 0.57857 0.52154]
[0.00265 0.98825 0.90528]]
[0.20764 0.29244]
[[0.51996 0.90184 0.98368 0.25752 0.56436 0.807 0.39437 0.73107]
[0.16096 0.60055 0.86596 0.98348 0.07938 0.42842 0.20453 0.45064]
[0.5476 0.0931 0.29701 0.92752 0.56902 0.45752 0.75351 0.74187]
[0.04849 0.70857 0.83933 0.1659 0.78101 0.2866 0.30646 0.66526]]
[0.11137 0.66482 0.88778 0.69627]
[[0.44029 0.43816 0.76502 0.56561]
[0.08485 0.5826 0.81474 0.33702]]
[0.92754 0.75067]
[[0.57379 0.7514 ]]
[0.07908]
```
基准测试
------------
1) 使用随机生成数据测试
```
./bench/dlrm_s_benchmark.sh
```
2) 使用[Criteo Kaggle Display Advertising Challenge Dataset](https://ailab.criteo.com/ressources/) 数据测试方法.
- 下载并解压数据到/data/kaggle路径下
```
mkdir -p /data/kaggle
tar xvf kaggle-display-advertising-challenge-dataset.tar.gz
```
- 执行测试脚本
```
./bench/dlrm_s_criteo_kaggle.sh [--test-freq=1024]
```
- 可以通过修改脚本中的以下参数来指定测试数据路径
- 首先可以指定训练数据地址 --raw-data-file=<path/train.txt>
- 可以指定预处理后的数据地址 --processed-data-file=<path/*.npz>
训练结果参考如下
<img src="./kaggle_dac_loss_accuracy_plots.png" width="900" height="320">
3) 多节点测试：代码支持分布式训练，目前支持gloo/nccl/mpi.
```
# 单节点4颗DCU测试，使用nccl通信，测试数据使用随机生成数据:
python -m torch.distributed.launch --nproc_per_node=8 dlrm_s_pytorch.py --arch-embedding-size="80000-80000-80000-80000-80000-80000-80000-80000" --arch-sparse-feature-size=64 --arch-mlp-bot="128-128-128-128" --arch-mlp-top="512-512-512-256-1" --max-ind-range=40000000
--data-generation=random --loss-function=bce --round-targets=True --learning-rate=1.0 --mini-batch-size=2048 --print-freq=2 --print-time --test-freq=2 --test-mini-batch-size=2048 --memory-map --use-gpu --num-batches=100 --dist-backend=nccl
# 多节点的情况可以添加如下参数:
--nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=1234
```
保存、加载模型参数
-------------------------------
* --save-model=<path/model.pt> ： 保存模型地址、名称
* --load-model=<path/model.pt> ： 加载模型
其他
----
想了解其他应用情况，可以参考地址：https://github.com/facebookresearch/dlrm
Version
-------
0.1 : Initial release of the DLRM code
1.0 : DLRM with distributed training, cpu support for row-wise adagrad optimizer
Requirements
------------
pytorch (*11/10/20*)
scikit-learn
numpy
onnx (*optional*)
pydot (*optional*)
torchviz (*optional*)
mpi (*optional for distributed backend*)
PyTorch/Recommendation/dlrm/bench/dlrm_s_benchmark.sh 0 → 100644
View file @ c43a53e4
#!/bin/bash
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#check if extra argument is passed to the test
if [[ $# == 1 ]]; then
dlrm_extra_option=$1
else
dlrm_extra_option=""
fi
#echo $dlrm_extra_option
cpu=1
gpu=1
pt=1
c2=1
ncores=28 #12 #6
nsockets="0"
ngpus="1 2 4 8"
numa_cmd="numactl --physcpubind=0-$((ncores-1)) -m $nsockets" #run on one socket, without HT
dlrm_pt_bin="python dlrm_s_pytorch.py"
dlrm_c2_bin="python dlrm_s_caffe2.py"
data=random #synthetic
print_freq=100
rand_seed=727
c2_net="async_scheduling"
#Model param
mb_size=2048 #1024 #512 #256
nbatches=1000 #500 #100
bot_mlp="512-512-64"
top_mlp="1024-1024-1024-1"
emb_size=64
nindices=100
emb="1000000-1000000-1000000-1000000-1000000-1000000-1000000-1000000"
interaction="dot"
tnworkers=0
tmb_size=16384
#_args="--mini-batch-size="${mb_size}\
_args=" --num-batches="${nbatches}\
" --data-generation="${data}\
" --arch-mlp-bot="${bot_mlp}\
" --arch-mlp-top="${top_mlp}\
" --arch-sparse-feature-size="${emb_size}\
" --arch-embedding-size="${emb}\
" --num-indices-per-lookup="${nindices}\
" --arch-interaction-op="${interaction}\
" --numpy-rand-seed="${rand_seed}\
" --print-freq="${print_freq}\
" --print-time"\
" --enable-profiling "
c2_args=" --caffe2-net-type="${c2_net}
# CPU Benchmarking
if [ $cpu = 1 ]; then
echo "--------------------------------------------"
echo "CPU Benchmarking - running on $ncores cores"
echo "--------------------------------------------"
if [ $pt = 1 ]; then
outf="model1_CPU_PT_$ncores.log"
outp="dlrm_s_pytorch.prof"
echo "-------------------------------"
echo "Running PT (log file: $outf)"
echo "-------------------------------"
cmd="$numa_cmd $dlrm_pt_bin --mini-batch-size=$mb_size --test-mini-batch-size=$tmb_size --test-num-workers=$tnworkers $_args $dlrm_extra_option > $outf"
echo $cmd
eval $cmd
min=$(grep "iteration" $outf | awk 'BEGIN{best=999999} {if (best > $7) best=$7} END{print best}')
echo "Min time per iteration = $min"
# move profiling file(s)
mv $outp ${outf//".log"/".prof"}
mv ${outp//".prof"/".json"} ${outf//".log"/".json"}
fi
if [ $c2 = 1 ]; then
outf="model1_CPU_C2_$ncores.log"
outp="dlrm_s_caffe2.prof"
echo "-------------------------------"
echo "Running C2 (log file: $outf)"
echo "-------------------------------"
cmd="$numa_cmd $dlrm_c2_bin --mini-batch-size=$mb_size $_args $c2_args $dlrm_extra_option 1> $outf 2> $outp"
echo $cmd
eval $cmd
min=$(grep "iteration" $outf | awk 'BEGIN{best=999999} {if (best > $7) best=$7} END{print best}')
echo "Min time per iteration = $min"
# move profiling file (collected from stderr above)
mv $outp ${outf//".log"/".prof"}
fi
fi
# GPU Benchmarking
if [ $gpu = 1 ]; then
echo "--------------------------------------------"
echo "GPU Benchmarking - running on $ngpus GPUs"
echo "--------------------------------------------"
for _ng in $ngpus
do
# weak scaling
# _mb_size=$((mb_size*_ng))
# strong scaling
_mb_size=$((mb_size*1))
_gpus=$(seq -s, 0 $((_ng-1)))
cuda_arg="CUDA_VISIBLE_DEVICES=$_gpus"
echo "-------------------"
echo "Using GPUS: "$_gpus
echo "-------------------"
if [ $pt = 1 ]; then
outf="model1_GPU_PT_$_ng.log"
outp="dlrm_s_pytorch.prof"
echo "-------------------------------"
echo "Running PT (log file: $outf)"
echo "-------------------------------"
cmd="$cuda_arg $dlrm_pt_bin --mini-batch-size=$_mb_size --test-mini-batch-size=$tmb_size --test-num-workers=$tnworkers $_args --use-gpu $dlrm_extra_option > $outf"
echo $cmd
eval $cmd
min=$(grep "iteration" $outf | awk 'BEGIN{best=999999} {if (best > $7) best=$7} END{print best}')
echo "Min time per iteration = $min"
# move profiling file(s)
mv $outp ${outf//".log"/".prof"}
mv ${outp//".prof"/".json"} ${outf//".log"/".json"}
fi
if [ $c2 = 1 ]; then
outf="model1_GPU_C2_$_ng.log"
outp="dlrm_s_caffe2.prof"
echo "-------------------------------"
echo "Running C2 (log file: $outf)"
echo "-------------------------------"
cmd="$cuda_arg $dlrm_c2_bin --mini-batch-size=$_mb_size $_args $c2_args --use-gpu $dlrm_extra_option 1> $outf 2> $outp"
echo $cmd
eval $cmd
min=$(grep "iteration" $outf | awk 'BEGIN{best=999999} {if (best > $7) best=$7} END{print best}')
echo "Min time per iteration = $min"
# move profiling file (collected from stderr above)
mv $outp ${outf//".log"/".prof"}
fi
done
fi
PyTorch/Recommendation/dlrm/bench/dlrm_s_criteo_kaggle.sh 0 → 100644
View file @ c43a53e4
#!/bin/bash
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
#WARNING: must have compiled PyTorch and caffe2
#check if extra argument is passed to the test
if [[ $# == 1 ]]; then
dlrm_extra_option=$1
else
dlrm_extra_option=""
fi
#echo $dlrm_extra_option
dlrm_pt_bin="python dlrm_s_pytorch.py"
#dlrm_c2_bin="python dlrm_s_caffe2.py"
echo "run pytorch ..."
# WARNING: the following parameters will be set based on the data set
# --arch-embedding-size=... (sparse feature sizes)
# --arch-mlp-bot=... (the input to the first layer of bottom mlp)
$dlrm_pt_bin --arch-sparse-feature-size=16 --arch-mlp-bot="13-512-256-64-16" --arch-mlp-top="512-256-1" --data-generation=dataset --data-set=kaggle --raw-data-file=/data/kaggle/train.txt --processed-data-file=/data/kaggle/kaggleAdDisplayChallenge_processed.npz --loss-function=bce --round-targets=True --learning-rate=0.1 --mini-batch-size=128 --print-freq=1024 --print-time --test-mini-batch-size=16384 --test-num-workers=16 $dlrm_extra_option 2>&1 | tee run_kaggle_pt.log
#echo "run caffe2 ..."
# WARNING: the following parameters will be set based on the data set
# --arch-embedding-size=... (sparse feature sizes)
# --arch-mlp-bot=... (the input to the first layer of bottom mlp)
#$dlrm_c2_bin --arch-sparse-feature-size=16 --arch-mlp-bot="13-512-256-64-16" --arch-mlp-top="512-256-1" --data-generation=dataset --data-set=kaggle --raw-data-file=./input/train.txt --processed-data-file=./input/kaggleAdDisplayChallenge_processed.npz --loss-function=bce --round-targets=True --learning-rate=0.1 --mini-batch-size=128 --print-freq=1024 --print-time $dlrm_extra_option 2>&1 | tee run_kaggle_c2.log
echo "done"
PyTorch/Recommendation/dlrm/bench/dlrm_s_criteo_terabyte.sh 0 → 100644
View file @ c43a53e4
#!/bin/bash
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
#WARNING: must have compiled PyTorch and caffe2
#check if extra argument is passed to the test
if [[ $# == 1 ]]; then
dlrm_extra_option=$1
else
dlrm_extra_option=""
fi
#echo $dlrm_extra_option
dlrm_pt_bin="python dlrm_s_pytorch.py"
dlrm_c2_bin="python dlrm_s_caffe2.py"
echo "run pytorch ..."
# WARNING: the following parameters will be set based on the data set
# --arch-embedding-size=... (sparse feature sizes)
# --arch-mlp-bot=... (the input to the first layer of bottom mlp)
$dlrm_pt_bin --arch-sparse-feature-size=64 --arch-mlp-bot="13-512-256-64" --arch-mlp-top="512-512-256-1" --max-ind-range=10000000 --data-generation=dataset --data-set=terabyte --raw-data-file=./input/day --processed-data-file=./input/terabyte_processed.npz --loss-function=bce --round-targets=True --learning-rate=0.1 --mini-batch-size=2048 --print-freq=1024 --print-time --test-mini-batch-size=16384 --test-num-workers=16 $dlrm_extra_option 2>&1 | tee run_terabyte_pt.log
echo "run caffe2 ..."
# WARNING: the following parameters will be set based on the data set
# --arch-embedding-size=... (sparse feature sizes)
# --arch-mlp-bot=... (the input to the first layer of bottom mlp)
$dlrm_c2_bin --arch-sparse-feature-size=64 --arch-mlp-bot="13-512-256-64" --arch-mlp-top="512-512-256-1" --max-ind-range=10000000 --data-generation=dataset --data-set=terabyte --raw-data-file=./input/day --processed-data-file=./input/terabyte_processed.npz --loss-function=bce --round-targets=True --learning-rate=0.1 --mini-batch-size=2048 --print-freq=1024 --print-time $dlrm_extra_option 2>&1 | tee run_terabyte_c2.log
echo "done"
PyTorch/Recommendation/dlrm/bench/run_and_time.sh 0 → 100644
View file @ c43a53e4
#!/bin/bash
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
#WARNING: must have compiled PyTorch and caffe2
#check if extra argument is passed to the test
if [[ $# == 1 ]]; then
dlrm_extra_option=$1
else
dlrm_extra_option=""
fi
#echo $dlrm_extra_option
python dlrm_s_pytorch.py --arch-sparse-feature-size=128 --arch-mlp-bot="13-512-256-128" --arch-mlp-top="1024-1024-512-256-1" --max-ind-range=40000000 --data-generation=dataset --data-set=terabyte --raw-data-file=./input/day --processed-data-file=./input/terabyte_processed.npz --loss-function=bce --round-targets=True --learning-rate=1.0 --mini-batch-size=2048 --print-freq=2048 --print-time --test-freq=102400 --test-mini-batch-size=16384 --test-num-workers=16 --memory-map --mlperf-logging --mlperf-auc-threshold=0.8025 --mlperf-bin-loader --mlperf-bin-shuffle $dlrm_extra_option 2>&1 | tee run_terabyte_mlperf_pt.log
echo "done"
PyTorch/Recommendation/dlrm/cython/cython_compile.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: compile .so from python code
from __future__ import absolute_import, division, print_function, unicode_literals
from setuptools import setup
from Cython.Build import cythonize
from distutils.extension import Extension
ext_modules = [
Extension(
"data_utils_cython",
["data_utils_cython.pyx"],
extra_compile_args=['-O3'],
extra_link_args=['-O3'],
)
]
setup(
name='data_utils_cython',
ext_modules=cythonize(ext_modules)
)
PyTorch/Recommendation/dlrm/cython/cython_criteo.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: run dataset pre-processing in standalone mode
# WARNING: These steps are required to work with Cython
# 1. Instal Cython
# > sudo yum install Cython
# 2. Please copy data_utils.py into data_utils_cython.pyx
# 3. Compile the data_utils_cython.pyx to generate .so
# (it's important to keep extension .pyx rather than .py
# to ensure the C/C++ .so no .py is loaded at import time)
# > python cython_compile.py build_ext --inplace
# This should create data_utils_cython.so, which can be loaded below with "import"
# 4. Run standalone datatset preprocessing to generate .npz files
# a. Kaggle
# > python cython_criteo.py --data-set=kaggle --raw-data-file=./input/train.txt
# --processed-data-file=./input/kaggleAdDisplayChallenge_processed.npz
# b. Terabyte
# > python cython_criteo.py --max-ind-range=10000000 [--memory-map] --data-set=terabyte
# --raw-data-file=./input/day --processed-data-file=./input/terabyte_processed.npz
from __future__ import absolute_import, division, print_function, unicode_literals
import data_utils_cython as duc
if __name__ == "__main__":
### import packages ###
import argparse
### parse arguments ###
parser = argparse.ArgumentParser(
description="Preprocess Criteo dataset"
)
# model related parameters
parser.add_argument("--max-ind-range", type=int, default=-1)
parser.add_argument("--data-sub-sample-rate", type=float, default=0.0) # in [0, 1]
parser.add_argument("--data-randomize", type=str, default="total") # or day or none
parser.add_argument("--memory-map", action="store_true", default=False)
parser.add_argument("--data-set", type=str, default="kaggle") # or terabyte
parser.add_argument("--raw-data-file", type=str, default="")
parser.add_argument("--processed-data-file", type=str, default="")
args = parser.parse_args()
duc.loadDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"train",
args.raw_data_file,
args.processed_data_file,
args.memory_map
)
PyTorch/Recommendation/dlrm/data_loader_terabyte.py 0 → 100644
View file @ c43a53e4
# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import absolute_import, division, print_function, unicode_literals
import os
import numpy as np
from torch.utils.data import Dataset
import torch
import time
import math
from tqdm import tqdm
import argparse
class DataLoader:
"""
DataLoader dedicated for the Criteo Terabyte Click Logs dataset
"""
def __init__(
self,
data_filename,
data_directory,
days,
batch_size,
max_ind_range=-1,
split="train",
drop_last_batch=False
):
self.data_filename = data_filename
self.data_directory = data_directory
self.days = days
self.batch_size = batch_size
self.max_ind_range = max_ind_range
total_file = os.path.join(
data_directory,
data_filename + "_day_count.npz"
)
with np.load(total_file) as data:
total_per_file = data["total_per_file"][np.array(days)]
self.length = sum(total_per_file)
if split == "test" or split == "val":
self.length = int(np.ceil(self.length / 2.))
self.split = split
self.drop_last_batch = drop_last_batch
def __iter__(self):
return iter(
_batch_generator(
self.data_filename, self.data_directory, self.days,
self.batch_size, self.split, self.drop_last_batch, self.max_ind_range
)
)
def __len__(self):
if self.drop_last_batch:
return self.length // self.batch_size
else:
return math.ceil(self.length / self.batch_size)
def _transform_features(
x_int_batch, x_cat_batch, y_batch, max_ind_range, flag_input_torch_tensor=False
):
if max_ind_range > 0:
x_cat_batch = x_cat_batch % max_ind_range
if flag_input_torch_tensor:
x_int_batch = torch.log(x_int_batch.clone().detach().type(torch.float) + 1)
x_cat_batch = x_cat_batch.clone().detach().type(torch.long)
y_batch = y_batch.clone().detach().type(torch.float32).view(-1, 1)
else:
x_int_batch = torch.log(torch.tensor(x_int_batch, dtype=torch.float) + 1)
x_cat_batch = torch.tensor(x_cat_batch, dtype=torch.long)
y_batch = torch.tensor(y_batch, dtype=torch.float32).view(-1, 1)
batch_size = x_cat_batch.shape[0]
feature_count = x_cat_batch.shape[1]
lS_o = torch.arange(batch_size).reshape(1, -1).repeat(feature_count, 1)
return x_int_batch, lS_o, x_cat_batch.t(), y_batch.view(-1, 1)
def _batch_generator(
data_filename, data_directory, days, batch_size, split, drop_last, max_ind_range
):
previous_file = None
for day in days:
filepath = os.path.join(
data_directory,
data_filename + "_{}_reordered.npz".format(day)
)
# print('Loading file: ', filepath)
with np.load(filepath) as data:
x_int = data["X_int"]
x_cat = data["X_cat"]
y = data["y"]
samples_in_file = y.shape[0]
batch_start_idx = 0
if split == "test" or split == "val":
length = int(np.ceil(samples_in_file / 2.))
if split == "test":
samples_in_file = length
elif split == "val":
batch_start_idx = samples_in_file - length
while batch_start_idx < samples_in_file - batch_size:
missing_samples = batch_size
if previous_file is not None:
missing_samples -= previous_file['y'].shape[0]
current_slice = slice(batch_start_idx, batch_start_idx + missing_samples)
x_int_batch = x_int[current_slice]
x_cat_batch = x_cat[current_slice]
y_batch = y[current_slice]
if previous_file is not None:
x_int_batch = np.concatenate(
[previous_file['x_int'], x_int_batch],
axis=0
)
x_cat_batch = np.concatenate(
[previous_file['x_cat'], x_cat_batch],
axis=0
)
y_batch = np.concatenate([previous_file['y'], y_batch], axis=0)
previous_file = None
if x_int_batch.shape[0] != batch_size:
raise ValueError('should not happen')
yield _transform_features(x_int_batch, x_cat_batch, y_batch, max_ind_range)
batch_start_idx += missing_samples
if batch_start_idx != samples_in_file:
current_slice = slice(batch_start_idx, samples_in_file)
if previous_file is not None:
previous_file = {
'x_int' : np.concatenate(
[previous_file['x_int'], x_int[current_slice]],
axis=0
),
'x_cat' : np.concatenate(
[previous_file['x_cat'], x_cat[current_slice]],
axis=0
),
'y' : np.concatenate([previous_file['y'], y[current_slice]], axis=0)
}
else:
previous_file = {
'x_int' : x_int[current_slice],
'x_cat' : x_cat[current_slice],
'y' : y[current_slice]
}
if not drop_last:
yield _transform_features(
previous_file['x_int'],
previous_file['x_cat'],
previous_file['y'],
max_ind_range
)
def _test():
generator = _batch_generator(
data_filename='day',
data_directory='./input',
days=range(23),
split="train",
batch_size=2048,
drop_last=True,
max_ind_range=-1
)
t1 = time.time()
for x_int, lS_o, x_cat, y in generator:
t2 = time.time()
time_diff = t2 - t1
t1 = t2
print(
"time {} x_int.shape: {} lS_o.shape: {} x_cat.shape: {} y.shape: {}".format(
time_diff, x_int.shape, lS_o.shape, x_cat.shape, y.shape
)
)
class CriteoBinDataset(Dataset):
"""Binary version of criteo dataset."""
def __init__(self, data_file, counts_file,
batch_size=1, max_ind_range=-1, bytes_per_feature=4):
# dataset
self.tar_fea = 1 # single target
self.den_fea = 13 # 13 dense features
self.spa_fea = 26 # 26 sparse features
self.tad_fea = self.tar_fea + self.den_fea
self.tot_fea = self.tad_fea + self.spa_fea
self.batch_size = batch_size
self.max_ind_range = max_ind_range
self.bytes_per_entry = (bytes_per_feature * self.tot_fea * batch_size)
self.num_entries = math.ceil(os.path.getsize(data_file) / self.bytes_per_entry)
print('data file:', data_file, 'number of batches:', self.num_entries)
self.file = open(data_file, 'rb')
with np.load(counts_file) as data:
self.counts = data["counts"]
# hardcoded for now
self.m_den = 13
def __len__(self):
return self.num_entries
def __getitem__(self, idx):
self.file.seek(idx * self.bytes_per_entry, 0)
raw_data = self.file.read(self.bytes_per_entry)
array = np.frombuffer(raw_data, dtype=np.int32)
tensor = torch.from_numpy(array).view((-1, self.tot_fea))
return _transform_features(x_int_batch=tensor[:, 1:14],
x_cat_batch=tensor[:, 14:],
y_batch=tensor[:, 0],
max_ind_range=self.max_ind_range,
flag_input_torch_tensor=True)
def __del__(self):
self.file.close()
def numpy_to_binary(input_files, output_file_path, split='train'):
"""Convert the data to a binary format to be read with CriteoBinDataset."""
# WARNING - both categorical and numerical data must fit into int32 for
# the following code to work correctly
with open(output_file_path, 'wb') as output_file:
if split == 'train':
for input_file in input_files:
print('Processing file: ', input_file)
np_data = np.load(input_file)
np_data = np.concatenate([np_data['y'].reshape(-1, 1),
np_data['X_int'],
np_data['X_cat']], axis=1)
np_data = np_data.astype(np.int32)
output_file.write(np_data.tobytes())
else:
assert len(input_files) == 1
np_data = np.load(input_files[0])
np_data = np.concatenate([np_data['y'].reshape(-1, 1),
np_data['X_int'],
np_data['X_cat']], axis=1)
np_data = np_data.astype(np.int32)
samples_in_file = np_data.shape[0]
midpoint = int(np.ceil(samples_in_file / 2.))
if split == "test":
begin = 0
end = midpoint
elif split == "val":
begin = midpoint
end = samples_in_file
else:
raise ValueError('Unknown split value: ', split)
output_file.write(np_data[begin:end].tobytes())
def _preprocess(args):
train_files = ['{}_{}_reordered.npz'.format(args.input_data_prefix, day) for
day in range(0, 23)]
test_valid_file = args.input_data_prefix + '_23_reordered.npz'
os.makedirs(args.output_directory, exist_ok=True)
for split in ['train', 'val', 'test']:
print('Running preprocessing for split =', split)
output_file = os.path.join(args.output_directory,
'{}_data.bin'.format(split))
input_files = train_files if split == 'train' else [test_valid_file]
numpy_to_binary(input_files=input_files,
output_file_path=output_file,
split=split)
def _test_bin():
parser = argparse.ArgumentParser()
parser.add_argument('--output_directory', required=True)
parser.add_argument('--input_data_prefix', required=True)
parser.add_argument('--split', choices=['train', 'test', 'val'],
required=True)
args = parser.parse_args()
_preprocess(args)
binary_data_file = os.path.join(args.output_directory,
'{}_data.bin'.format(args.split))
counts_file = os.path.join(args.output_directory, 'day_fea_count.npz')
dataset_binary = CriteoBinDataset(data_file=binary_data_file,
counts_file=counts_file,
batch_size=2048,)
from dlrm_data_pytorch import CriteoDataset
from dlrm_data_pytorch import collate_wrapper_criteo_offset as collate_wrapper_criteo
binary_loader = torch.utils.data.DataLoader(
dataset_binary,
batch_size=None,
shuffle=False,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
)
original_dataset = CriteoDataset(
dataset='terabyte',
max_ind_range=10 * 1000 * 1000,
sub_sample_rate=1,
randomize=True,
split=args.split,
raw_path=args.input_data_prefix,
pro_data='dummy_string',
memory_map=True
)
original_loader = torch.utils.data.DataLoader(
original_dataset,
batch_size=2048,
shuffle=False,
num_workers=0,
collate_fn=collate_wrapper_criteo,
pin_memory=False,
drop_last=False,
)
assert len(dataset_binary) == len(original_loader)
for i, (old_batch, new_batch) in tqdm(enumerate(zip(original_loader,
binary_loader)),
total=len(dataset_binary)):
for j in range(len(new_batch)):
if not np.array_equal(old_batch[j], new_batch[j]):
raise ValueError('FAILED: Datasets not equal')
if i > len(dataset_binary):
break
print('PASSED')
if __name__ == '__main__':
_test()
_test_bin()
PyTorch/Recommendation/dlrm/data_utils.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: generate inputs and targets for the DLRM benchmark
#
# Utility function(s) to download and pre-process public data sets
# - Criteo Kaggle Display Advertising Challenge Dataset
# https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset
# - Criteo Terabyte Dataset
# https://labs.criteo.com/2013/12/download-terabyte-click-logs
#
# After downloading dataset, run:
# getCriteoAdData(
# datafile="<path-to-train.txt>",
# o_filename=kaggleAdDisplayChallenge_processed.npz,
# max_ind_range=-1,
# sub_sample_rate=0.0,
# days=7,
# data_split='train',
# randomize='total',
# criteo_kaggle=True,
# memory_map=False
# )
# getCriteoAdData(
# datafile="<path-to-day_{0,...,23}>",
# o_filename=terabyte_processed.npz,
# max_ind_range=-1,
# sub_sample_rate=0.0,
# days=24,
# data_split='train',
# randomize='total',
# criteo_kaggle=False,
# memory_map=False
# )
from __future__ import absolute_import, division, print_function, unicode_literals
import sys
# import os
from os import path
from multiprocessing import Process, Manager
# import io
# from io import StringIO
# import collections as coll
import numpy as np
def convertUStringToDistinctIntsDict(mat, convertDicts, counts):
# Converts matrix of unicode strings into distinct integers.
#
# Inputs:
# mat (np.array): array of unicode strings to convert
# convertDicts (list): dictionary for each column
# counts (list): number of different categories in each column
#
# Outputs:
# out (np.array): array of output integers
# convertDicts (list): dictionary for each column
# counts (list): number of different categories in each column
# check if convertDicts and counts match correct length of mat
if len(convertDicts) != mat.shape[1] or len(counts) != mat.shape[1]:
print("Length of convertDicts or counts does not match input shape")
print("Generating convertDicts and counts...")
convertDicts = [{} for _ in range(mat.shape[1])]
counts = [0 for _ in range(mat.shape[1])]
# initialize output
out = np.zeros(mat.shape)
for j in range(mat.shape[1]):
for i in range(mat.shape[0]):
# add to convertDict and increment count
if mat[i, j] not in convertDicts[j]:
convertDicts[j][mat[i, j]] = counts[j]
counts[j] += 1
out[i, j] = convertDicts[j][mat[i, j]]
return out, convertDicts, counts
def convertUStringToDistinctIntsUnique(mat, mat_uni, counts):
# mat is an array of 0,...,# samples, with each being 26 categorical features
# check if mat_unique and counts match correct length of mat
if len(mat_uni) != mat.shape[1] or len(counts) != mat.shape[1]:
print("Length of mat_unique or counts does not match input shape")
print("Generating mat_unique and counts...")
mat_uni = [np.array([]) for _ in range(mat.shape[1])]
counts = [0 for _ in range(mat.shape[1])]
# initialize output
out = np.zeros(mat.shape)
ind_map = [np.array([]) for _ in range(mat.shape[1])]
# find out and assign unique ids to features
for j in range(mat.shape[1]):
m = mat_uni[j].size
mat_concat = np.concatenate((mat_uni[j], mat[:, j]))
mat_uni[j], ind_map[j] = np.unique(mat_concat, return_inverse=True)
out[:, j] = ind_map[j][m:]
counts[j] = mat_uni[j].size
return out, mat_uni, counts
def processCriteoAdData(d_path, d_file, npzfile, i, convertDicts, pre_comp_counts):
# Process Kaggle Display Advertising Challenge or Terabyte Dataset
# by converting unicode strings in X_cat to integers and
# converting negative integer values in X_int.
#
# Loads data in the form "{kaggle|terabyte}_day_i.npz" where i is the day.
#
# Inputs:
# d_path (str): path for {kaggle|terabyte}_day_i.npz files
# i (int): splits in the dataset (typically 0 to 7 or 0 to 24)
# process data if not all files exist
filename_i = npzfile + "_{0}_processed.npz".format(i)
if path.exists(filename_i):
print("Using existing " + filename_i, end="\n")
else:
print("Not existing " + filename_i)
with np.load(npzfile + "_{0}.npz".format(i)) as data:
# categorical features
'''
# Approach 1a: using empty dictionaries
X_cat, convertDicts, counts = convertUStringToDistinctIntsDict(
data["X_cat"], convertDicts, counts
)
'''
'''
# Approach 1b: using empty np.unique
X_cat, convertDicts, counts = convertUStringToDistinctIntsUnique(
data["X_cat"], convertDicts, counts
)
'''
# Approach 2a: using pre-computed dictionaries
X_cat_t = np.zeros(data["X_cat_t"].shape)
for j in range(26):
for k, x in enumerate(data["X_cat_t"][j, :]):
X_cat_t[j, k] = convertDicts[j][x]
# continuous features
X_int = data["X_int"]
X_int[X_int < 0] = 0
# targets
y = data["y"]
np.savez_compressed(
filename_i,
# X_cat = X_cat,
X_cat=np.transpose(X_cat_t), # transpose of the data
X_int=X_int,
y=y,
)
print("Processed " + filename_i, end="\n")
# sanity check (applicable only if counts have been pre-computed & are re-computed)
# for j in range(26):
# if pre_comp_counts[j] != counts[j]:
# sys.exit("ERROR: Sanity check on counts has failed")
# print("\nSanity check on counts passed")
return
def concatCriteoAdData(
d_path,
d_file,
npzfile,
trafile,
days,
data_split,
randomize,
total_per_file,
total_count,
memory_map,
o_filename
):
# Concatenates different days and saves the result.
#
# Inputs:
# days (int): total number of days in the dataset (typically 7 or 24)
# d_path (str): path for {kaggle|terabyte}_day_i.npz files
# o_filename (str): output file name
#
# Output:
# o_file (str): output file path
if memory_map:
# dataset break up per fea
# tar_fea = 1 # single target
den_fea = 13 # 13 dense features
spa_fea = 26 # 26 sparse features
# tad_fea = tar_fea + den_fea
# tot_fea = tad_fea + spa_fea
# create offset per file
offset_per_file = np.array([0] + [x for x in total_per_file])
for i in range(days):
offset_per_file[i + 1] += offset_per_file[i]
'''
# Approach 1, 2 and 3 use indices, while Approach 4 does not use them
# create indices
indices = np.arange(total_count)
if data_split == "none":
if randomize == "total":
indices = np.random.permutation(indices)
else:
indices = np.array_split(indices, offset_per_file[1:-1])
# randomize train data (per day)
if randomize == "day": # or randomize == "total":
for i in range(len(indices) - 1):
indices[i] = np.random.permutation(indices[i])
print("Randomized indices per day ...")
train_indices = np.concatenate(indices[:-1])
test_indices = indices[-1]
# randomize train data (across days)
if randomize == "total":
train_indices = np.random.permutation(train_indices)
print("Randomized indices across days ...")
indices = np.concatenate((train_indices, test_indices))
# no reordering
# indices = np.arange(total_count)
'''
'''
# Approach 1: simple and slow (no grouping is used)
# check if data already exists
recreate_flag = False
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered.npy".format(j)
if path.exists(filename_j):
print("Using existing " + filename_j)
else:
recreate_flag = True
# load, reorder and concatenate data (memmap all reordered files per feature)
if recreate_flag:
# init reordered files (.npy appended automatically)
z = np.zeros((total_count))
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered".format(j)
np.save(filename_j, z)
print("Creating " + filename_j)
for i in range(days):
filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
with np.load(filename_i) as data:
X_cat_t = np.transpose(data["X_cat"])
X_int_t = np.transpose(data["X_int"])
y = data["y"]
size = len(y)
# sanity check
if total_per_file[i] != size:
sys.exit("ERROR: sanity check on number of samples failed")
# setup start and end ranges
start = offset_per_file[i]
end = offset_per_file[i + 1]
# print(filename_i)
# print("start=" + str(start) + " end=" + str(end)
# + " diff=" + str(end - start) + "=" + str(total_per_file[i]))
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered.npy".format(j)
fj = np.load(filename_j, mmap_mode='r+')
if j < tar_fea:
fj[indices[start:end]] = y
elif tar_fea <= j and j < tad_fea:
fj[indices[start:end]] = X_int_t[j - tar_fea, :]
else:
fj[indices[start:end]] = X_cat_t[j - tad_fea, :]
del fj
else:
print("Reordered fea files already exist, skipping ...")
# check if data already exists
recreate_flag = False
for i in range(days):
filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
if path.exists(filename_i):
print("Using existing " + filename_i)
else:
recreate_flag = True
# split reordered data by files (memmap all reordered files per feature)
# on the day boundary del the file object and memmap again
if recreate_flag:
for i in range(days):
filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
size = total_per_file[i]
X_int_t = np.zeros((den_fea, size))
X_cat_t = np.zeros((spa_fea, size))
# setup start and end ranges
start = offset_per_file[i]
end = offset_per_file[i + 1]
print("Creating " + filename_i)
# print("start=" + str(start) + " end=" + str(end)
# + " diff=" + str(end - start) + "=" + str(total_per_file[i]))
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered.npy".format(j)
fj = np.load(filename_j, mmap_mode='r')
if j < tar_fea:
y = fj[start:end]
elif tar_fea <= j and j < tad_fea:
X_int_t[j - tar_fea, :] = fj[start:end]
else:
X_cat_t[j - tad_fea, :] = fj[start:end]
del fj
np.savez_compressed(
filename_i,
X_cat=np.transpose(X_cat_t), # transpose of the data
X_int=np.transpose(X_int_t), # transpose of the data
y=y,
)
else:
print("Reordered day files already exist, skipping ...")
'''
'''
# Approach 2: group days
# check if data already exists
recreate_flag = False
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered.npy".format(j)
if path.exists(filename_j):
print("Using existing " + filename_j)
else:
recreate_flag = True
# load, reorder and concatenate data (memmap all reordered files per feature)
if recreate_flag:
# init reordered files (.npy appended automatically)
z = np.zeros((total_count))
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered".format(j)
np.save(filename_j, z)
print("Creating " + filename_j)
group_day = 3 # e.g. 8, 4 or 3
group_num = days // group_day
file_group = [i*group_day for i in range(group_num)] + [days]
for ii in range(group_num):
# for last may be group_size != group_num, therefore reset it below
group_size = file_group[ii + 1] - file_group[ii]
X_cat_t = [0]*group_size
X_int_t = [0]*group_size
y = [0]*group_size
start = [0]*group_size
end = [0]*group_size
for ig in range(group_size):
i = file_group[ii] + ig
filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
# setup start and end ranges
start[ig] = offset_per_file[i]
end[ig] = offset_per_file[i + 1]
# print(filename_i)
# load a group of files
with np.load(filename_i) as data:
X_cat_t[ig] = np.transpose(data["X_cat"])
X_int_t[ig] = np.transpose(data["X_int"])
y[ig] = data["y"]
# sanity check
if total_per_file[i] != len(y[ig]):
sys.exit("ERROR: sanity check on number of samples failed")
# print("start=" + str(start) + " end=" + str(end)
# + " diff=" + str(end[ig]-start[ig]) + "=" + str(total_per_file[i]))
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered.npy".format(j)
fj = np.load(filename_j, mmap_mode='r+')
for ig in range(group_size):
if j < tar_fea:
fj[indices[start[ig]:end[ig]]] = y[ig]
elif tar_fea <= j and j < tad_fea:
fj[indices[start[ig]:end[ig]]] = X_int_t[ig][j - tar_fea, :]
else:
fj[indices[start[ig]:end[ig]]] = X_cat_t[ig][j - tad_fea, :]
del fj
else:
print("Reordered fea files already exist, skipping ...")
# check if data already exists
recreate_flag = False
for i in range(days):
filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
if path.exists(filename_i):
print("Using existing " + filename_i)
else:
recreate_flag = True
# split reordered data by files (memmap all reordered files per feature)
# on the day boundary del the file object and memmap again
if recreate_flag:
for ii in range(group_num):
# for last may be group_size != group_num, therefore reset it below
group_size = file_group[ii + 1] - file_group[ii]
X_cat_t= []; X_int_t = []
for ig in range(group_size):
i = file_group[ii] + ig
X_int_t.append(np.zeros((den_fea, total_per_file[i])))
X_cat_t.append(np.zeros((spa_fea, total_per_file[i])))
y = [0]*group_size
start = [0]*group_size
end = [0]*group_size
for j in range(tot_fea):
filename_j = trafile + "_{0}_reordered.npy".format(j)
fj = np.load(filename_j, mmap_mode='r')
# load a group of files
for ig in range(group_size):
i = file_group[ii] + ig
# setup start and end ranges
start[ig] = offset_per_file[i]
end[ig] = offset_per_file[i + 1]
# load data for the group of files
if j < tar_fea:
y[ig] = fj[start[ig]:end[ig]]
elif tar_fea <= j and j < tad_fea:
X_int_t[ig][j - tar_fea, :] = fj[start[ig]:end[ig]]
else:
X_cat_t[ig][j - tad_fea, :] = fj[start[ig]:end[ig]]
del fj
for ig in range(group_size):
i = file_group[ii] + ig
filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
print("Creating " + filename_i)
np.savez_compressed(
filename_i,
X_cat=np.transpose(X_cat_t[ig]), # transpose of the data
X_int=np.transpose(X_int_t[ig]), # transpose of the data
y=y[ig],
)
else:
print("Reordered day files already exist, skipping ...")
'''
'''
# Approach 3: group features
# check if data already exists
group_fea = 5 # e.g. 8, 5 or 4
group_num = tot_fea // group_fea
if tot_fea % group_fea != 0: # sanity check
sys.exit("ERROR: the group_fea must divided tot_fea evenly.")
recreate_flag = False
for jn in range(group_num):
filename_j = trafile + "_{0}_reordered{1}.npy".format(
jn, group_fea
)
if path.exists(filename_j):
print("Using existing " + filename_j)
else:
recreate_flag = True
# load, reorder and concatenate data (memmap all reordered files per feature)
if recreate_flag:
# init reordered files (.npy appended automatically)
z = np.zeros((group_fea, total_count))
for jn in range(group_num):
filename_j = trafile + "_{0}_reordered{1}".format(
jn, group_fea
)
np.save(filename_j, z)
print("Creating " + filename_j)
for i in range(days):
filename_i = d_path + npzfile + "_{0}_processed.npz".format(i)
with np.load(filename_i) as data:
X_cat_t = np.transpose(data["X_cat"])
X_int_t = np.transpose(data["X_int"])
y = data["y"]
size = len(y)
# sanity check
if total_per_file[i] != size:
sys.exit("ERROR: sanity check on number of samples failed")
# setup start and end ranges
start = offset_per_file[i]
end = offset_per_file[i + 1]
# print(filename_i)
# print("start=" + str(start) + " end=" + str(end)
# + " diff=" + str(end - start) + "=" + str(total_per_file[i]))
for jn in range(group_num):
filename_j = trafile + "_{0}_reordered{1}.npy".format(
jn, group_fea
)
fj = np.load(filename_j, mmap_mode='r+')
for jg in range(group_fea):
j = jn * group_fea + jg
# print("j=" + str(j) + " jn=" + str(jn) + " jg=" + str(jg))
if j < tar_fea:
fj[jg, indices[start:end]] = y
elif tar_fea <= j and j < tad_fea:
fj[jg, indices[start:end]] = X_int_t[j - tar_fea, :]
else:
fj[jg, indices[start:end]] = X_cat_t[j - tad_fea, :]
del fj
else:
print("Reordered fea files already exist, skipping ...")
# check if data already exists
recreate_flag = False
for i in range(days):
filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
if path.exists(filename_i):
print("Using existing" + filename_i)
else:
recreate_flag = True
# split reordered data by files (memmap all reordered files per feature)
# on the day boundary del the file object and memmap again
if recreate_flag:
for i in range(days):
filename_i = d_path + npzfile + "_{0}_reordered.npz".format(i)
size = total_per_file[i]
X_int_t = np.zeros((den_fea, size))
X_cat_t = np.zeros((spa_fea, size))
# setup start and end ranges
start = offset_per_file[i]
end = offset_per_file[i + 1]
print("Creating " + filename_i)
# print("start=" + str(start) + " end=" + str(end)
# + " diff=" + str(end - start) + "=" + str(total_per_file[i]))
for jn in range(group_num):
filename_j = trafile + "_{0}_reordered{1}.npy".format(
jn, group_fea
)
fj = np.load(filename_j, mmap_mode='r')
for jg in range(group_fea):
j = jn * group_fea + jg
# print("j=" + str(j) + " jn=" + str(jn) + " jg=" + str(jg))
if j < tar_fea:
y = fj[jg, start:end]
elif tar_fea <= j and j < tad_fea:
X_int_t[j - tar_fea, :] = fj[jg, start:end]
else:
X_cat_t[j - tad_fea, :] = fj[jg, start:end]
del fj
np.savez_compressed(
filename_i,
X_cat=np.transpose(X_cat_t), # transpose of the data
X_int=np.transpose(X_int_t), # transpose of the data
y=y,
)
else:
print("Reordered day files already exist, skipping ...")
'''
# Approach 4: Fisher-Yates-Rao (FYR) shuffle algorithm
# 1st pass of FYR shuffle
# check if data already exists
recreate_flag = False
for j in range(days):
filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
if (
path.exists(filename_j_y)
and path.exists(filename_j_d)
and path.exists(filename_j_s)
):
print(
"Using existing\n"
+ filename_j_y + "\n"
+ filename_j_d + "\n"
+ filename_j_s
)
else:
recreate_flag = True
# reorder across buckets using sampling
if recreate_flag:
# init intermediate files (.npy appended automatically)
for j in range(days):
filename_j_y = npzfile + "_{0}_intermediate_y".format(j)
filename_j_d = npzfile + "_{0}_intermediate_d".format(j)
filename_j_s = npzfile + "_{0}_intermediate_s".format(j)
np.save(filename_j_y, np.zeros((total_per_file[j])))
np.save(filename_j_d, np.zeros((total_per_file[j], den_fea)))
np.save(filename_j_s, np.zeros((total_per_file[j], spa_fea)))
# start processing files
total_counter = [0] * days
for i in range(days):
filename_i = npzfile + "_{0}_processed.npz".format(i)
with np.load(filename_i) as data:
X_cat = data["X_cat"]
X_int = data["X_int"]
y = data["y"]
size = len(y)
# sanity check
if total_per_file[i] != size:
sys.exit("ERROR: sanity check on number of samples failed")
# debug prints
print("Reordering (1st pass) " + filename_i)
# create buckets using sampling of random ints
# from (discrete) uniform distribution
buckets = []
for _j in range(days):
buckets.append([])
counter = [0] * days
days_to_sample = days if data_split == "none" else days - 1
if randomize == "total":
rand_u = np.random.randint(low=0, high=days_to_sample, size=size)
for k in range(size):
# sample and make sure elements per buckets do not overflow
if data_split == "none" or i < days - 1:
# choose bucket
p = rand_u[k]
# retry of the bucket is full
while total_counter[p] + counter[p] >= total_per_file[p]:
p = np.random.randint(low=0, high=days_to_sample)
else: # preserve the last day/bucket if needed
p = i
buckets[p].append(k)
counter[p] += 1
else: # randomize is day or none
for k in range(size):
# do not sample, preserve the data in this bucket
p = i
buckets[p].append(k)
counter[p] += 1
# sanity check
if np.sum(counter) != size:
sys.exit("ERROR: sanity check on number of samples failed")
# debug prints
# print(counter)
# print(str(np.sum(counter)) + " = " + str(size))
# print([len(x) for x in buckets])
# print(total_counter)
# partially feel the buckets
for j in range(days):
filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
start = total_counter[j]
end = total_counter[j] + counter[j]
# target buckets
fj_y = np.load(filename_j_y, mmap_mode='r+')
# print("start=" + str(start) + " end=" + str(end)
# + " end - start=" + str(end - start) + " "
# + str(fj_y[start:end].shape) + " "
# + str(len(buckets[j])))
fj_y[start:end] = y[buckets[j]]
del fj_y
# dense buckets
fj_d = np.load(filename_j_d, mmap_mode='r+')
# print("start=" + str(start) + " end=" + str(end)
# + " end - start=" + str(end - start) + " "
# + str(fj_d[start:end, :].shape) + " "
# + str(len(buckets[j])))
fj_d[start:end, :] = X_int[buckets[j], :]
del fj_d
# sparse buckets
fj_s = np.load(filename_j_s, mmap_mode='r+')
# print("start=" + str(start) + " end=" + str(end)
# + " end - start=" + str(end - start) + " "
# + str(fj_s[start:end, :].shape) + " "
# + str(len(buckets[j])))
fj_s[start:end, :] = X_cat[buckets[j], :]
del fj_s
# update counters for next step
total_counter[j] += counter[j]
# 2nd pass of FYR shuffle
# check if data already exists
for j in range(days):
filename_j = npzfile + "_{0}_reordered.npz".format(j)
if path.exists(filename_j):
print("Using existing " + filename_j)
else:
recreate_flag = True
# reorder within buckets
if recreate_flag:
for j in range(days):
filename_j_y = npzfile + "_{0}_intermediate_y.npy".format(j)
filename_j_d = npzfile + "_{0}_intermediate_d.npy".format(j)
filename_j_s = npzfile + "_{0}_intermediate_s.npy".format(j)
fj_y = np.load(filename_j_y)
fj_d = np.load(filename_j_d)
fj_s = np.load(filename_j_s)
indices = range(total_per_file[j])
if randomize == "day" or randomize == "total":
if data_split == "none" or j < days - 1:
indices = np.random.permutation(range(total_per_file[j]))
filename_r = npzfile + "_{0}_reordered.npz".format(j)
print("Reordering (2nd pass) " + filename_r)
np.savez_compressed(
filename_r,
X_cat=fj_s[indices, :],
X_int=fj_d[indices, :],
y=fj_y[indices],
)
'''
# sanity check (under no reordering norms should be zero)
for i in range(days):
filename_i_o = npzfile + "_{0}_processed.npz".format(i)
print(filename_i_o)
with np.load(filename_i_o) as data_original:
X_cat_o = data_original["X_cat"]
X_int_o = data_original["X_int"]
y_o = data_original["y"]
filename_i_r = npzfile + "_{0}_reordered.npz".format(i)
print(filename_i_r)
with np.load(filename_i_r) as data_reordered:
X_cat_r = data_reordered["X_cat"]
X_int_r = data_reordered["X_int"]
y_r = data_reordered["y"]
print(np.linalg.norm(y_o - y_r))
print(np.linalg.norm(X_int_o - X_int_r))
print(np.linalg.norm(X_cat_o - X_cat_r))
'''
else:
print("Concatenating multiple days into %s.npz file" % str(d_path + o_filename))
# load and concatenate data
for i in range(days):
filename_i = npzfile + "_{0}_processed.npz".format(i)
with np.load(filename_i) as data:
if i == 0:
X_cat = data["X_cat"]
X_int = data["X_int"]
y = data["y"]
else:
X_cat = np.concatenate((X_cat, data["X_cat"]))
X_int = np.concatenate((X_int, data["X_int"]))
y = np.concatenate((y, data["y"]))
print("Loaded day:", i, "y = 1:", len(y[y == 1]), "y = 0:", len(y[y == 0]))
with np.load(d_path + d_file + "_fea_count.npz") as data:
counts = data["counts"]
print("Loaded counts!")
np.savez_compressed(
d_path + o_filename + ".npz",
X_cat=X_cat,
X_int=X_int,
y=y,
counts=counts,
)
return d_path + o_filename + ".npz"
def transformCriteoAdData(X_cat, X_int, y, days, data_split, randomize, total_per_file):
# Transforms Criteo Kaggle or terabyte data by applying log transformation
# on dense features and converting everything to appropriate tensors.
#
# Inputs:
# X_cat (ndarray): array of integers corresponding to preprocessed
# categorical features
# X_int (ndarray): array of integers corresponding to dense features
# y (ndarray): array of bool corresponding to labels
# data_split(str): flag for splitting dataset into training/validation/test
# sets
# randomize (str): determines randomization scheme
# "none": no randomization
# "day": randomizes each day"s data (only works if split = True)
# "total": randomizes total dataset
#
# Outputs:
# if split:
# X_cat_train (tensor): sparse features for training set
# X_int_train (tensor): dense features for training set
# y_train (tensor): labels for training set
# X_cat_val (tensor): sparse features for validation set
# X_int_val (tensor): dense features for validation set
# y_val (tensor): labels for validation set
# X_cat_test (tensor): sparse features for test set
# X_int_test (tensor): dense features for test set
# y_test (tensor): labels for test set
# else:
# X_cat (tensor): sparse features
# X_int (tensor): dense features
# y (tensor): label
# define initial set of indices
indices = np.arange(len(y))
# create offset per file
offset_per_file = np.array([0] + [x for x in total_per_file])
for i in range(days):
offset_per_file[i + 1] += offset_per_file[i]
# split dataset
if data_split == 'train':
indices = np.array_split(indices, offset_per_file[1:-1])
# randomize train data (per day)
if randomize == "day": # or randomize == "total":
for i in range(len(indices) - 1):
indices[i] = np.random.permutation(indices[i])
print("Randomized indices per day ...")
train_indices = np.concatenate(indices[:-1])
test_indices = indices[-1]
test_indices, val_indices = np.array_split(test_indices, 2)
print("Defined training and testing indices...")
# randomize train data (across days)
if randomize == "total":
train_indices = np.random.permutation(train_indices)
print("Randomized indices across days ...")
# indices = np.concatenate((train_indices, test_indices))
# create training, validation, and test sets
X_cat_train = X_cat[train_indices]
X_int_train = X_int[train_indices]
y_train = y[train_indices]
X_cat_val = X_cat[val_indices]
X_int_val = X_int[val_indices]
y_val = y[val_indices]
X_cat_test = X_cat[test_indices]
X_int_test = X_int[test_indices]
y_test = y[test_indices]
print("Split data according to indices...")
X_cat_train = X_cat_train.astype(np.long)
X_int_train = np.log(X_int_train.astype(np.float32) + 1)
y_train = y_train.astype(np.float32)
X_cat_val = X_cat_val.astype(np.long)
X_int_val = np.log(X_int_val.astype(np.float32) + 1)
y_val = y_val.astype(np.float32)
X_cat_test = X_cat_test.astype(np.long)
X_int_test = np.log(X_int_test.astype(np.float32) + 1)
y_test = y_test.astype(np.float32)
print("Converted to tensors...done!")
return (
X_cat_train,
X_int_train,
y_train,
X_cat_val,
X_int_val,
y_val,
X_cat_test,
X_int_test,
y_test,
)
else:
# randomize data
if randomize == "total":
indices = np.random.permutation(indices)
print("Randomized indices...")
X_cat = X_cat[indices].astype(np.long)
X_int = np.log(X_int[indices].astype(np.float32) + 1)
y = y[indices].astype(np.float32)
print("Converted to tensors...done!")
return (X_cat, X_int, y, [], [], [], [], [], [])
def getCriteoAdData(
datafile,
o_filename,
max_ind_range=-1,
sub_sample_rate=0.0,
days=7,
data_split='train',
randomize='total',
criteo_kaggle=True,
memory_map=False,
dataset_multiprocessing=False,
):
# Passes through entire dataset and defines dictionaries for categorical
# features and determines the number of total categories.
#
# Inputs:
# datafile : path to downloaded raw data file
# o_filename (str): saves results under o_filename if filename is not ""
#
# Output:
# o_file (str): output file path
#split the datafile into path and filename
lstr = datafile.split("/")
d_path = "/".join(lstr[0:-1]) + "/"
d_file = lstr[-1].split(".")[0] if criteo_kaggle else lstr[-1]
npzfile = d_path + ((d_file + "_day") if criteo_kaggle else d_file)
trafile = d_path + ((d_file + "_fea") if criteo_kaggle else "fea")
# count number of datapoints in training set
total_file = d_path + d_file + "_day_count.npz"
if path.exists(total_file):
with np.load(total_file) as data:
total_per_file = list(data["total_per_file"])
total_count = np.sum(total_per_file)
print("Skipping counts per file (already exist)")
else:
total_count = 0
total_per_file = []
if criteo_kaggle:
# WARNING: The raw data consists of a single train.txt file
# Each line in the file is a sample, consisting of 13 continuous and
# 26 categorical features (an extra space indicates that feature is
# missing and will be interpreted as 0).
if path.exists(datafile):
print("Reading data from path=%s" % (datafile))
with open(str(datafile)) as f:
for _ in f:
total_count += 1
total_per_file.append(total_count)
# reset total per file due to split
num_data_per_split, extras = divmod(total_count, days)
total_per_file = [num_data_per_split] * days
for j in range(extras):
total_per_file[j] += 1
# split into days (simplifies code later on)
file_id = 0
boundary = total_per_file[file_id]
nf = open(npzfile + "_" + str(file_id), "w")
with open(str(datafile)) as f:
for j, line in enumerate(f):
if j == boundary:
nf.close()
file_id += 1
nf = open(npzfile + "_" + str(file_id), "w")
boundary += total_per_file[file_id]
nf.write(line)
nf.close()
else:
sys.exit("ERROR: Criteo Kaggle Display Ad Challenge Dataset path is invalid; please download from https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset")
else:
# WARNING: The raw data consist of day_0.gz,... ,day_23.gz text files
# Each line in the file is a sample, consisting of 13 continuous and
# 26 categorical features (an extra space indicates that feature is
# missing and will be interpreted as 0).
for i in range(days):
datafile_i = datafile + "_" + str(i) # + ".gz"
if path.exists(str(datafile_i)):
print("Reading data from path=%s" % (str(datafile_i)))
# file day_<number>
total_per_file_count = 0
with open(str(datafile_i)) as f:
for _ in f:
total_per_file_count += 1
total_per_file.append(total_per_file_count)
total_count += total_per_file_count
else:
sys.exit("ERROR: Criteo Terabyte Dataset path is invalid; please download from https://labs.criteo.com/2013/12/download-terabyte-click-logs")
# process a file worth of data and reinitialize data
# note that a file main contain a single or multiple splits
def process_one_file(
datfile,
npzfile,
split,
num_data_in_split,
dataset_multiprocessing,
convertDictsDay=None,
resultDay=None
):
if dataset_multiprocessing:
convertDicts_day = [{} for _ in range(26)]
with open(str(datfile)) as f:
y = np.zeros(num_data_in_split, dtype="i4") # 4 byte int
X_int = np.zeros((num_data_in_split, 13), dtype="i4") # 4 byte int
X_cat = np.zeros((num_data_in_split, 26), dtype="i4") # 4 byte int
if sub_sample_rate == 0.0:
rand_u = 1.0
else:
rand_u = np.random.uniform(low=0.0, high=1.0, size=num_data_in_split)
i = 0
percent = 0
for k, line in enumerate(f):
# process a line (data point)
line = line.split('\t')
# set missing values to zero
for j in range(len(line)):
if (line[j] == '') or (line[j] == '\n'):
line[j] = '0'
# sub-sample data by dropping zero targets, if needed
target = np.int32(line[0])
if target == 0 and \
(rand_u if sub_sample_rate == 0.0 else rand_u[k]) < sub_sample_rate:
continue
y[i] = target
X_int[i] = np.array(line[1:14], dtype=np.int32)
if max_ind_range > 0:
X_cat[i] = np.array(
list(map(lambda x: int(x, 16) % max_ind_range, line[14:])),
dtype=np.int32
)
else:
X_cat[i] = np.array(
list(map(lambda x: int(x, 16), line[14:])),
dtype=np.int32
)
# count uniques
if dataset_multiprocessing:
for j in range(26):
convertDicts_day[j][X_cat[i][j]] = 1
# debug prints
if float(i)/num_data_in_split*100 > percent+1:
percent = int(float(i)/num_data_in_split*100)
print(
"Load %d/%d (%d%%) Split: %d Label True: %d Stored: %d"
% (
i,
num_data_in_split,
percent,
split,
target,
y[i],
),
end="\n",
)
else:
for j in range(26):
convertDicts[j][X_cat[i][j]] = 1
# debug prints
print(
"Load %d/%d Split: %d Label True: %d Stored: %d"
% (
i,
num_data_in_split,
split,
target,
y[i],
),
end="\r",
)
i += 1
# store num_data_in_split samples or extras at the end of file
# count uniques
# X_cat_t = np.transpose(X_cat)
# for j in range(26):
# for x in X_cat_t[j,:]:
# convertDicts[j][x] = 1
# store parsed
filename_s = npzfile + "_{0}.npz".format(split)
if path.exists(filename_s):
print("\nSkip existing " + filename_s)
else:
np.savez_compressed(
filename_s,
X_int=X_int[0:i, :],
# X_cat=X_cat[0:i, :],
X_cat_t=np.transpose(X_cat[0:i, :]), # transpose of the data
y=y[0:i],
)
print("\nSaved " + npzfile + "_{0}.npz!".format(split))
if dataset_multiprocessing:
resultDay[split] = i
convertDictsDay[split] = convertDicts_day
return
else:
return i
# create all splits (reuse existing files if possible)
recreate_flag = False
convertDicts = [{} for _ in range(26)]
# WARNING: to get reproducable sub-sampling results you must reset the seed below
# np.random.seed(123)
# in this case there is a single split in each day
for i in range(days):
npzfile_i = npzfile + "_{0}.npz".format(i)
npzfile_p = npzfile + "_{0}_processed.npz".format(i)
if path.exists(npzfile_i):
print("Skip existing " + npzfile_i)
elif path.exists(npzfile_p):
print("Skip existing " + npzfile_p)
else:
recreate_flag = True
if recreate_flag:
if dataset_multiprocessing:
resultDay = Manager().dict()
convertDictsDay = Manager().dict()
processes = [Process(target=process_one_file,
name="process_one_file:%i" % i,
args=(npzfile + "_{0}".format(i),
npzfile,
i,
total_per_file[i],
dataset_multiprocessing,
convertDictsDay,
resultDay,
)
) for i in range(0, days)]
for process in processes:
process.start()
for process in processes:
process.join()
for day in range(days):
total_per_file[day] = resultDay[day]
print("Constructing convertDicts Split: {}".format(day))
convertDicts_tmp = convertDictsDay[day]
for i in range(26):
for j in convertDicts_tmp[i]:
convertDicts[i][j] = 1
else:
for i in range(days):
total_per_file[i] = process_one_file(
npzfile + "_{0}".format(i),
npzfile,
i,
total_per_file[i],
dataset_multiprocessing,
)
# report and save total into a file
total_count = np.sum(total_per_file)
if not path.exists(total_file):
np.savez_compressed(total_file, total_per_file=total_per_file)
print("Total number of samples:", total_count)
print("Divided into days/splits:\n", total_per_file)
# dictionary files
counts = np.zeros(26, dtype=np.int32)
if recreate_flag:
# create dictionaries
for j in range(26):
for i, x in enumerate(convertDicts[j]):
convertDicts[j][x] = i
dict_file_j = d_path + d_file + "_fea_dict_{0}.npz".format(j)
if not path.exists(dict_file_j):
np.savez_compressed(
dict_file_j,
unique=np.array(list(convertDicts[j]), dtype=np.int32)
)
counts[j] = len(convertDicts[j])
# store (uniques and) counts
count_file = d_path + d_file + "_fea_count.npz"
if not path.exists(count_file):
np.savez_compressed(count_file, counts=counts)
else:
# create dictionaries (from existing files)
for j in range(26):
with np.load(d_path + d_file + "_fea_dict_{0}.npz".format(j)) as data:
unique = data["unique"]
for i, x in enumerate(unique):
convertDicts[j][x] = i
# load (uniques and) counts
with np.load(d_path + d_file + "_fea_count.npz") as data:
counts = data["counts"]
# process all splits
if dataset_multiprocessing:
processes = [Process(target=processCriteoAdData,
name="processCriteoAdData:%i" % i,
args=(d_path,
d_file,
npzfile,
i,
convertDicts,
counts,
)
) for i in range(0, days)]
for process in processes:
process.start()
for process in processes:
process.join()
else:
for i in range(days):
processCriteoAdData(d_path, d_file, npzfile, i, convertDicts, counts)
o_file = concatCriteoAdData(
d_path,
d_file,
npzfile,
trafile,
days,
data_split,
randomize,
total_per_file,
total_count,
memory_map,
o_filename
)
return o_file
def loadDataset(
dataset,
max_ind_range,
sub_sample_rate,
randomize,
data_split,
raw_path="",
pro_data="",
memory_map=False
):
# dataset
if dataset == "kaggle":
days = 7
o_filename = "kaggleAdDisplayChallenge_processed"
elif dataset == "terabyte":
days = 24
o_filename = "terabyte_processed"
else:
raise(ValueError("Data set option is not supported"))
# split the datafile into path and filename
lstr = raw_path.split("/")
d_path = "/".join(lstr[0:-1]) + "/"
d_file = lstr[-1].split(".")[0] if dataset == "kaggle" else lstr[-1]
npzfile = (d_file + "_day") if dataset == "kaggle" else d_file
# trafile = d_path + ((d_file + "_fea") if dataset == "kaggle" else "fea")
# check if pre-processed data is available
data_ready = True
if memory_map:
for i in range(days):
reo_data = d_path + npzfile + "_{0}_reordered.npz".format(i)
if not path.exists(str(reo_data)):
data_ready = False
else:
if not path.exists(str(pro_data)):
data_ready = False
# pre-process data if needed
# WARNNING: when memory mapping is used we get a collection of files
if data_ready:
print("Reading pre-processed data=%s" % (str(pro_data)))
file = str(pro_data)
else:
print("Reading raw data=%s" % (str(raw_path)))
file = getCriteoAdData(
raw_path,
o_filename,
max_ind_range,
sub_sample_rate,
days,
data_split,
randomize,
dataset == "kaggle",
memory_map
)
return file, days
if __name__ == "__main__":
### import packages ###
import argparse
### parse arguments ###
parser = argparse.ArgumentParser(
description="Preprocess Criteo dataset"
)
# model related parameters
parser.add_argument("--max-ind-range", type=int, default=-1)
parser.add_argument("--data-sub-sample-rate", type=float, default=0.0) # in [0, 1]
parser.add_argument("--data-randomize", type=str, default="total") # or day or none
parser.add_argument("--memory-map", action="store_true", default=False)
parser.add_argument("--data-set", type=str, default="kaggle") # or terabyte
parser.add_argument("--raw-data-file", type=str, default="")
parser.add_argument("--processed-data-file", type=str, default="")
args = parser.parse_args()
loadDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"train",
args.raw_data_file,
args.processed_data_file,
args.memory_map
)
PyTorch/Recommendation/dlrm/dlrm_data_caffe2.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: generate inputs and targets for the dlrm benchmark
# The inpts and outputs are generated according to the following three option(s)
# 1) random distribution
# 2) synthetic distribution, based on unique accesses and distances between them
# i) R. Hassan, A. Harris, N. Topham and A. Efthymiou "Synthetic Trace-Driven
# Simulation of Cache Memory", IEEE AINAM'07
# 3) public data set
# i) Criteo Kaggle Display Advertising Challenge Dataset
# https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset
# ii) Criteo Terabyte Dataset
# https://labs.criteo.com/2013/12/download-terabyte-click-logs
from __future__ import absolute_import, division, print_function, unicode_literals
import bisect
import collections
# others
# from os import path
import sys
import data_utils
# numpy
import numpy as np
# pytorch
import torch
from numpy import random as ra
from torch.utils.data import Dataset
# Kaggle Display Advertising Challenge Dataset
# dataset (str): name of dataset (Kaggle or Terabyte)
# randomize (str): determines randomization scheme
# 'none': no randomization
# 'day': randomizes each day's data (only works if split = True)
# 'total': randomizes total dataset
# split (bool) : to split into train, test, validation data-sets
class CriteoDatasetWMemoryMap(Dataset):
def __init__(
self,
dataset,
max_ind_range,
sub_sample_rate,
randomize,
split="train",
raw_path="",
pro_data="",
):
# dataset
# tar_fea = 1 # single target
den_fea = 13 # 13 dense features
# spa_fea = 26 # 26 sparse features
# tad_fea = tar_fea + den_fea
# tot_fea = tad_fea + spa_fea
if dataset == "kaggle":
days = 7
elif dataset == "terabyte":
days = 24
else:
raise (ValueError("Data set option is not supported"))
self.max_ind_range = max_ind_range
# split the datafile into path and filename
lstr = raw_path.split("/")
self.d_path = "/".join(lstr[0:-1]) + "/"
self.d_file = lstr[-1].split(".")[0] if dataset == "kaggle" else lstr[-1]
self.npzfile = self.d_path + (
(self.d_file + "_day") if dataset == "kaggle" else self.d_file
)
self.trafile = self.d_path + (
(self.d_file + "_fea") if dataset == "kaggle" else "fea"
)
# get a number of samples per day
total_file = self.d_path + self.d_file + "_day_count.npz"
with np.load(total_file) as data:
total_per_file = data["total_per_file"]
# compute offsets per file
self.offset_per_file = np.array([0] + list(total_per_file))
for i in range(days):
self.offset_per_file[i + 1] += self.offset_per_file[i]
# print(self.offset_per_file)
# setup data
self.split = split
if split == "none" or split == "train":
self.day = 0
self.max_day_range = days if split == "none" else days - 1
elif split == "test" or split == "val":
self.day = days - 1
num_samples = self.offset_per_file[days] - self.offset_per_file[days - 1]
self.test_size = int(np.ceil(num_samples / 2.0))
self.val_size = num_samples - self.test_size
else:
sys.exit("ERROR: dataset split is neither none, nor train or test.")
# load unique counts
with np.load(self.d_path + self.d_file + "_fea_count.npz") as data:
self.counts = data["counts"]
self.m_den = den_fea # X_int.shape[1]
self.n_emb = len(self.counts)
print("Sparse features= %d, Dense features= %d" % (self.n_emb, self.m_den))
# Load the test data
# Only a single day is used for testing
if self.split == "test" or self.split == "val":
# only a single day is used for testing
fi = self.npzfile + "_{0}_reordered.npz".format(self.day)
with np.load(fi) as data:
self.X_int = data["X_int"] # continuous feature
self.X_cat = data["X_cat"] # categorical feature
self.y = data["y"] # target
def __getitem__(self, index):
if isinstance(index, slice):
return [
self[idx]
for idx in range(
index.start or 0, index.stop or len(self), index.step or 1
)
]
if self.split == "none" or self.split == "train":
# check if need to swicth to next day and load data
if index == self.offset_per_file[self.day]:
# print("day_boundary switch", index)
self.day_boundary = self.offset_per_file[self.day]
fi = self.npzfile + "_{0}_reordered.npz".format(self.day)
# print('Loading file: ', fi)
with np.load(fi) as data:
self.X_int = data["X_int"] # continuous feature
self.X_cat = data["X_cat"] # categorical feature
self.y = data["y"] # target
self.day = (self.day + 1) % self.max_day_range
i = index - self.day_boundary
elif self.split == "test" or self.split == "val":
# only a single day is used for testing
i = index + (0 if self.split == "test" else self.test_size)
else:
sys.exit("ERROR: dataset split is neither none, nor train or test.")
if self.max_ind_range > 0:
return self.X_int[i], self.X_cat[i] % self.max_ind_range, self.y[i]
else:
return self.X_int[i], self.X_cat[i], self.y[i]
def _default_preprocess(self, X_int, X_cat, y):
X_int = torch.log(torch.tensor(X_int, dtype=torch.float) + 1)
if self.max_ind_range > 0:
X_cat = torch.tensor(X_cat % self.max_ind_range, dtype=torch.long)
else:
X_cat = torch.tensor(X_cat, dtype=torch.long)
y = torch.tensor(y.astype(np.float32))
return X_int, X_cat, y
def __len__(self):
if self.split == "none":
return self.offset_per_file[-1]
elif self.split == "train":
return self.offset_per_file[-2]
elif self.split == "test":
return self.test_size
elif self.split == "val":
return self.val_size
else:
sys.exit("ERROR: dataset split is neither none, nor train nor test.")
def collate_wrapper_criteo(list_of_tuples):
# where each tuple is (X_int, X_cat, y)
transposed_data = list(zip(*list_of_tuples))
X_int = torch.log(torch.tensor(transposed_data[0], dtype=torch.float) + 1)
X_cat = torch.tensor(transposed_data[1], dtype=torch.long)
T = torch.tensor(transposed_data[2], dtype=torch.float32).view(-1, 1)
batchSize = X_cat.shape[0]
featureCnt = X_cat.shape[1]
lS_i = [X_cat[:, i] for i in range(featureCnt)]
lS_o = [torch.tensor(range(batchSize)) for _ in range(featureCnt)]
return X_int, torch.stack(lS_o), torch.stack(lS_i), T
# Conversion from offset to length
def offset_to_length_convertor(lS_o, lS_i):
def diff(tensor):
return tensor[1:] - tensor[:-1]
return torch.stack(
[
diff(torch.cat((S_o, torch.tensor(lS_i[ind].shape))).int())
for ind, S_o in enumerate(lS_o)
]
)
def unpack_batch(b, data_gen, data_set):
return b[0], b[1], b[2], b[3], torch.ones(b[3].size())
def read_dataset(
dataset,
max_ind_range,
sub_sample_rate,
mini_batch_size,
num_batches,
randomize,
split="train",
raw_data="",
processed_data="",
memory_map=False,
inference_only=False,
test_mini_batch_size=1,
):
# split the datafile into path and filename
lstr = raw_data.split("/")
d_path = "/".join(lstr[0:-1]) + "/"
d_file = lstr[-1].split(".")[0] if dataset == "kaggle" else lstr[-1]
# npzfile = d_path + ((d_file + "_day") if dataset == "kaggle" else d_file)
# trafile = d_path + ((d_file + "_fea") if dataset == "kaggle" else "fea")
# load
print("Loading %s dataset..." % dataset)
nbatches = 0
file, days = data_utils.loadDataset(
dataset,
max_ind_range,
sub_sample_rate,
randomize,
split,
raw_data,
processed_data,
memory_map,
)
if memory_map:
# WARNING: at this point the data has been reordered and shuffled across files
# e.g. day_<number>_reordered.npz, what remains is simply to read and feed
# the data from each file, going in the order of days file-by-file, to the
# model during training.
train_data = CriteoDatasetWMemoryMap(
dataset,
max_ind_range,
sub_sample_rate,
randomize,
"train",
raw_data,
processed_data,
)
test_data = CriteoDatasetWMemoryMap(
dataset,
max_ind_range,
sub_sample_rate,
randomize,
"test",
raw_data,
processed_data,
)
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=mini_batch_size,
shuffle=False,
num_workers=0,
collate_fn=collate_wrapper_criteo,
pin_memory=False,
drop_last=False, # True
)
test_loader = torch.utils.data.DataLoader(
test_data,
batch_size=test_mini_batch_size,
shuffle=False,
num_workers=0,
collate_fn=collate_wrapper_criteo,
pin_memory=False,
drop_last=False, # True
)
return train_data, train_loader, test_data, test_loader
else:
# load and preprocess data
with np.load(file) as data:
X_int = data["X_int"]
X_cat = data["X_cat"]
y = data["y"]
counts = data["counts"]
# get a number of samples per day
total_file = d_path + d_file + "_day_count.npz"
with np.load(total_file) as data:
total_per_file = data["total_per_file"]
# transform
(
X_cat_train,
X_int_train,
y_train,
X_cat_val,
X_int_val,
y_val,
X_cat_test,
X_int_test,
y_test,
) = data_utils.transformCriteoAdData(
X_cat, X_int, y, days, split, randomize, total_per_file
)
ln_emb = counts
m_den = X_int_train.shape[1]
n_emb = len(counts)
print("Sparse features = %d, Dense features = %d" % (n_emb, m_den))
# adjust parameters
def assemble_samples(X_cat, X_int, y, max_ind_range, print_message):
if max_ind_range > 0:
X_cat = X_cat % max_ind_range
nsamples = len(y)
data_size = nsamples
# using floor is equivalent to dropping last mini-batch (drop_last = True)
nbatches = int(np.floor((data_size * 1.0) / mini_batch_size))
print(print_message)
if num_batches != 0 and num_batches < nbatches:
print(
"Limiting to %d batches of the total % d batches"
% (num_batches, nbatches)
)
nbatches = num_batches
else:
print("Total number of batches %d" % nbatches)
# data main loop
lX = []
lS_lengths = []
lS_indices = []
lT = []
for j in range(0, nbatches):
# number of data points in a batch
print("Reading in batch: %d / %d" % (j + 1, nbatches), end="\r")
n = min(mini_batch_size, data_size - (j * mini_batch_size))
# dense feature
idx_start = j * mini_batch_size
lX.append((X_int[idx_start : (idx_start + n)]).astype(np.float32))
# Targets - outputs
lT.append(
(y[idx_start : idx_start + n]).reshape(-1, 1).astype(np.int32)
)
# sparse feature (sparse indices)
lS_emb_indices = []
# for each embedding generate a list of n lookups,
# where each lookup is composed of multiple sparse indices
for size in range(n_emb):
lS_batch_indices = []
for _b in range(n):
# num of sparse indices to be used per embedding, e.g. for
# store lengths and indices
lS_batch_indices += (
(X_cat[idx_start + _b][size].reshape(-1)).astype(np.int32)
).tolist()
lS_emb_indices.append(lS_batch_indices)
lS_indices.append(lS_emb_indices)
# Criteo Kaggle data it is 1 because data is categorical
lS_lengths.append(
[(list(np.ones(n).astype(np.int32))) for _ in range(n_emb)]
)
print("\n")
return nbatches, lX, lS_lengths, lS_indices, lT
# adjust training data
(nbatches, lX, lS_lengths, lS_indices, lT) = assemble_samples(
X_cat_train, X_int_train, y_train, max_ind_range, "Training data"
)
# adjust testing data
(nbatches_t, lX_t, lS_lengths_t, lS_indices_t, lT_t) = assemble_samples(
X_cat_test, X_int_test, y_test, max_ind_range, "Testing data"
)
# end if memory_map
return (
nbatches,
lX,
lS_lengths,
lS_indices,
lT,
nbatches_t,
lX_t,
lS_lengths_t,
lS_indices_t,
lT_t,
ln_emb,
m_den,
)
def generate_random_data(
m_den,
ln_emb,
data_size,
num_batches,
mini_batch_size,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
num_targets=1,
round_targets=False,
data_generation="random",
trace_file="",
enable_padding=False,
):
nbatches = int(np.ceil((data_size * 1.0) / mini_batch_size))
if num_batches != 0:
nbatches = num_batches
data_size = nbatches * mini_batch_size
# print("Total number of batches %d" % nbatches)
# inputs and targets
lT = []
lX = []
lS_lengths = []
lS_indices = []
for j in range(0, nbatches):
# number of data points in a batch
n = min(mini_batch_size, data_size - (j * mini_batch_size))
# generate a batch of dense and sparse features
if data_generation == "random":
(Xt, lS_emb_lengths, lS_emb_indices) = generate_uniform_input_batch(
m_den, ln_emb, n, num_indices_per_lookup, num_indices_per_lookup_fixed
)
elif data_generation == "synthetic":
(Xt, lS_emb_lengths, lS_emb_indices) = generate_synthetic_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
trace_file,
enable_padding,
)
else:
sys.exit(
"ERROR: --data-generation=" + data_generation + " is not supported"
)
# dense feature
lX.append(Xt)
# sparse feature (sparse indices)
lS_lengths.append(lS_emb_lengths)
lS_indices.append(lS_emb_indices)
# generate a batch of target (probability of a click)
P = generate_random_output_batch(n, num_targets, round_targets)
lT.append(P)
return (nbatches, lX, lS_lengths, lS_indices, lT)
def generate_random_output_batch(n, num_targets=1, round_targets=False):
# target (probability of a click)
if round_targets:
P = np.round(ra.rand(n, num_targets).astype(np.float32)).astype(np.int32)
else:
P = ra.rand(n, num_targets).astype(np.float32)
return P
# uniform ditribution (input data)
def generate_uniform_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
):
# dense feature
Xt = ra.rand(n, m_den).astype(np.float32)
# sparse feature (sparse indices)
lS_emb_lengths = []
lS_emb_indices = []
# for each embedding generate a list of n lookups,
# where each lookup is composed of multiple sparse indices
for size in ln_emb:
lS_batch_lengths = []
lS_batch_indices = []
for _ in range(n):
# num of sparse indices to be used per embedding (between
if num_indices_per_lookup_fixed:
sparse_group_size = np.int32(num_indices_per_lookup)
else:
# random between [1,num_indices_per_lookup])
r = ra.random(1)
sparse_group_size = np.int32(
max(1, np.round(r * min(size, num_indices_per_lookup))[0])
)
# sparse indices to be used per embedding
r = ra.random(sparse_group_size)
sparse_group = np.unique(np.round(r * (size - 1)).astype(np.int32))
# reset sparse_group_size in case some index duplicates were removed
sparse_group_size = np.int32(sparse_group.size)
# store lengths and indices
lS_batch_lengths += [sparse_group_size]
lS_batch_indices += sparse_group.tolist()
lS_emb_lengths.append(lS_batch_lengths)
lS_emb_indices.append(lS_batch_indices)
return (Xt, lS_emb_lengths, lS_emb_indices)
# synthetic distribution (input data)
def generate_synthetic_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
trace_file,
enable_padding=False,
):
# dense feature
Xt = ra.rand(n, m_den).astype(np.float32)
# sparse feature (sparse indices)
lS_emb_lengths = []
lS_emb_indices = []
# for each embedding generate a list of n lookups,
# where each lookup is composed of multiple sparse indices
for i, size in enumerate(ln_emb):
lS_batch_lengths = []
lS_batch_indices = []
for _ in range(n):
# num of sparse indices to be used per embedding (between
if num_indices_per_lookup_fixed:
sparse_group_size = np.int32(num_indices_per_lookup)
else:
# random between [1,num_indices_per_lookup])
r = ra.random(1)
sparse_group_size = np.int32(
max(1, np.round(r * min(size, num_indices_per_lookup))[0])
)
# sparse indices to be used per embedding
file_path = trace_file
line_accesses, list_sd, cumm_sd = read_dist_from_file(
file_path.replace("j", str(i))
)
# debug print
# print('input')
# print(line_accesses); print(list_sd); print(cumm_sd);
# print(sparse_group_size)
# approach 1: rand
# r = trace_generate_rand(
# line_accesses, list_sd, cumm_sd, sparse_group_size, enable_padding
# )
# approach 2: lru
r = trace_generate_lru(
line_accesses, list_sd, cumm_sd, sparse_group_size, enable_padding
)
# WARNING: if the distribution in the file is not consistent with
# embedding table dimensions, below mod guards against out of
# range access
sparse_group = np.unique(r).astype(np.int32)
minsg = np.min(sparse_group)
maxsg = np.max(sparse_group)
if (minsg < 0) or (size <= maxsg):
print(
"WARNING: distribution is inconsistent with embedding "
+ "table size (using mod to recover and continue)"
)
sparse_group = np.mod(sparse_group, size).astype(np.int32)
# sparse_group = np.unique(np.array(np.mod(r, size-1)).astype(np.int32))
# reset sparse_group_size in case some index duplicates were removed
sparse_group_size = np.int32(sparse_group.size)
# store lengths and indices
lS_batch_lengths += [sparse_group_size]
lS_batch_indices += sparse_group.tolist()
lS_emb_lengths.append(lS_batch_lengths)
lS_emb_indices.append(lS_batch_indices)
return (Xt, lS_emb_lengths, lS_emb_indices)
def generate_stack_distance(cumm_val, cumm_dist, max_i, i, enable_padding=False):
u = ra.rand(1)
if i < max_i:
# only generate stack distances up to the number of new references seen so far
j = bisect.bisect(cumm_val, i) - 1
fi = cumm_dist[j]
u *= fi # shrink distribution support to exclude last values
elif enable_padding:
# WARNING: disable generation of new references (once all have been seen)
fi = cumm_dist[0]
u = (1.0 - fi) * u + fi # remap distribution support to exclude first value
for (j, f) in enumerate(cumm_dist):
if u <= f:
return cumm_val[j]
# WARNING: global define, must be consistent across all synthetic functions
cache_line_size = 1
def trace_generate_lru(
line_accesses, list_sd, cumm_sd, out_trace_len, enable_padding=False
):
max_sd = list_sd[-1]
l = len(line_accesses)
i = 0
ztrace = []
for _ in range(out_trace_len):
sd = generate_stack_distance(list_sd, cumm_sd, max_sd, i, enable_padding)
mem_ref_within_line = 0 # floor(ra.rand(1)*cache_line_size) #0
# generate memory reference
if sd == 0: # new reference #
line_ref = line_accesses.pop(0)
line_accesses.append(line_ref)
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
i += 1
else: # existing reference #
line_ref = line_accesses[l - sd]
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
line_accesses.pop(l - sd)
line_accesses.append(line_ref)
# save generated memory reference
ztrace.append(mem_ref)
return ztrace
def trace_generate_rand(
line_accesses, list_sd, cumm_sd, out_trace_len, enable_padding=False
):
max_sd = list_sd[-1]
l = len(line_accesses) # !!!Unique,
i = 0
ztrace = []
for _ in range(out_trace_len):
sd = generate_stack_distance(list_sd, cumm_sd, max_sd, i, enable_padding)
mem_ref_within_line = 0 # floor(ra.rand(1)*cache_line_size) #0
# generate memory reference
if sd == 0: # new reference #
line_ref = line_accesses.pop(0)
line_accesses.append(line_ref)
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
i += 1
else: # existing reference #
line_ref = line_accesses[l - sd]
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
ztrace.append(mem_ref)
return ztrace
def trace_profile(trace, enable_padding=False):
# number of elements in the array (assuming 1D)
# n = trace.size
rstack = [] # S
stack_distances = [] # SDS
line_accesses = [] # L
for x in trace:
r = np.uint64(x / cache_line_size)
l = len(rstack)
try: # found #
i = rstack.index(r)
# WARNING: I believe below is the correct depth in terms of meaning of the
# algorithm, but that is not what seems to be in the paper alg.
# -1 can be subtracted if we defined the distance between
# consecutive accesses (e.g. r, r) as 0 rather than 1.
sd = l - i # - 1
# push r to the end of stack_distances
stack_distances.insert(0, sd)
# remove r from its position and insert to the top of stack
rstack.pop(i) # rstack.remove(r)
rstack.insert(l - 1, r)
except ValueError: # not found #
sd = 0 # -1
# push r to the end of stack_distances/line_accesses
stack_distances.insert(0, sd)
line_accesses.insert(0, r)
# push r to the top of stack
rstack.insert(l, r)
if enable_padding:
# WARNING: notice that as the ratio between the number of samples (l)
# and cardinality (c) of a sample increases the probability of
# generating a sample gets smaller and smaller because there are
# few new samples compared to repeated samples. This means that for a
# long trace with relatively small cardinality it will take longer to
# generate all new samples and therefore obtain full distribution support
# and hence it takes longer for distribution to resemble the original.
# Therefore, we may pad the number of new samples to be on par with
# average number of samples l/c artificially.
l = len(stack_distances)
c = max(stack_distances)
padding = int(np.ceil(l / c))
stack_distances = stack_distances + [0] * padding
return (rstack, stack_distances, line_accesses)
# auxiliary read/write routines
def read_trace_from_file(file_path):
try:
with open(file_path) as f:
if args.trace_file_binary_type:
array = np.fromfile(f, dtype=np.uint64)
trace = array.astype(np.uint64).tolist()
else:
line = f.readline()
trace = list(map(lambda x: np.uint64(x), line.split(", ")))
return trace
except Exception:
print("ERROR: no input trace file has been provided")
def write_trace_to_file(file_path, trace):
try:
if args.trace_file_binary_type:
with open(file_path, "wb+") as f:
np.array(trace).astype(np.uint64).tofile(f)
else:
with open(file_path, "w+") as f:
s = str(trace)
f.write(s[1 : len(s) - 1])
except Exception:
print("ERROR: no output trace file has been provided")
def read_dist_from_file(file_path):
try:
with open(file_path, "r") as f:
lines = f.read().splitlines()
except Exception:
print("Wrong file or file path")
# read unique accesses
unique_accesses = [int(el) for el in lines[0].split(", ")]
# read cumulative distribution (elements are passed as two separate lists)
list_sd = [int(el) for el in lines[1].split(", ")]
cumm_sd = [float(el) for el in lines[2].split(", ")]
return unique_accesses, list_sd, cumm_sd
def write_dist_to_file(file_path, unique_accesses, list_sd, cumm_sd):
try:
with open(file_path, "w") as f:
# unique_acesses
s = str(unique_accesses)
f.write(s[1 : len(s) - 1] + "\n")
# list_sd
s = str(list_sd)
f.write(s[1 : len(s) - 1] + "\n")
# cumm_sd
s = str(cumm_sd)
f.write(s[1 : len(s) - 1] + "\n")
except Exception:
print("Wrong file or file path")
if __name__ == "__main__":
import sys
import operator
import argparse
### parse arguments ###
parser = argparse.ArgumentParser(description="Generate Synthetic Distributions")
parser.add_argument("--trace-file", type=str, default="./input/trace.log")
parser.add_argument("--trace-file-binary-type", type=bool, default=False)
parser.add_argument("--trace-enable-padding", type=bool, default=False)
parser.add_argument("--dist-file", type=str, default="./input/dist.log")
parser.add_argument(
"--synthetic-file", type=str, default="./input/trace_synthetic.log"
)
parser.add_argument("--numpy-rand-seed", type=int, default=123)
parser.add_argument("--print-precision", type=int, default=5)
args = parser.parse_args()
### some basic setup ###
np.random.seed(args.numpy_rand_seed)
np.set_printoptions(precision=args.print_precision)
### read trace ###
trace = read_trace_from_file(args.trace_file)
# print(trace)
### profile trace ###
(_, stack_distances, line_accesses) = trace_profile(
trace, args.trace_enable_padding
)
stack_distances.reverse()
line_accesses.reverse()
# print(line_accesses)
# print(stack_distances)
### compute probability distribution ###
# count items
l = len(stack_distances)
dc = sorted(
collections.Counter(stack_distances).items(), key=operator.itemgetter(0)
)
# create a distribution
list_sd = list(map(lambda tuple_x_k: tuple_x_k[0], dc)) # x = tuple_x_k[0]
dist_sd = list(
map(lambda tuple_x_k: tuple_x_k[1] / float(l), dc)
) # k = tuple_x_k[1]
cumm_sd = [] # np.cumsum(dc).tolist() #prefixsum
for i, (_, k) in enumerate(dc):
if i == 0:
cumm_sd.append(k / float(l))
else:
# add the 2nd element of the i-th tuple in the dist_sd list
cumm_sd.append(cumm_sd[i - 1] + (k / float(l)))
### write stack_distance and line_accesses to a file ###
write_dist_to_file(args.dist_file, line_accesses, list_sd, cumm_sd)
### generate correspondinf synthetic ###
# line_accesses, list_sd, cumm_sd = read_dist_from_file(args.dist_file)
synthetic_trace = trace_generate_lru(
line_accesses, list_sd, cumm_sd, len(trace), args.trace_enable_padding
)
# synthetic_trace = trace_generate_rand(
# line_accesses, list_sd, cumm_sd, len(trace), args.trace_enable_padding
# )
write_trace_to_file(args.synthetic_file, synthetic_trace)
PyTorch/Recommendation/dlrm/dlrm_data_pytorch.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: generate inputs and targets for the dlrm benchmark
# The inpts and outputs are generated according to the following three option(s)
# 1) random distribution
# 2) synthetic distribution, based on unique accesses and distances between them
# i) R. Hassan, A. Harris, N. Topham and A. Efthymiou "Synthetic Trace-Driven
# Simulation of Cache Memory", IEEE AINAM'07
# 3) public data set
# i) Criteo Kaggle Display Advertising Challenge Dataset
# https://labs.criteo.com/2014/02/kaggle-display-advertising-challenge-dataset
# ii) Criteo Terabyte Dataset
# https://labs.criteo.com/2013/12/download-terabyte-click-logs
from __future__ import absolute_import, division, print_function, unicode_literals
# others
from os import path
import sys
import bisect
import collections
import data_utils
# numpy
import numpy as np
from numpy import random as ra
from collections import deque
# pytorch
import torch
from torch.utils.data import Dataset, RandomSampler
import data_loader_terabyte
import mlperf_logger
# Kaggle Display Advertising Challenge Dataset
# dataset (str): name of dataset (Kaggle or Terabyte)
# randomize (str): determines randomization scheme
# "none": no randomization
# "day": randomizes each day"s data (only works if split = True)
# "total": randomizes total dataset
# split (bool) : to split into train, test, validation data-sets
class CriteoDataset(Dataset):
def __init__(
self,
dataset,
max_ind_range,
sub_sample_rate,
randomize,
split="train",
raw_path="",
pro_data="",
memory_map=False,
dataset_multiprocessing=False,
):
# dataset
# tar_fea = 1 # single target
den_fea = 13 # 13 dense features
# spa_fea = 26 # 26 sparse features
# tad_fea = tar_fea + den_fea
# tot_fea = tad_fea + spa_fea
if dataset == "kaggle":
days = 7
out_file = "kaggleAdDisplayChallenge_processed"
elif dataset == "terabyte":
days = 24
out_file = "terabyte_processed"
else:
raise(ValueError("Data set option is not supported"))
self.max_ind_range = max_ind_range
self.memory_map = memory_map
# split the datafile into path and filename
lstr = raw_path.split("/")
self.d_path = "/".join(lstr[0:-1]) + "/"
self.d_file = lstr[-1].split(".")[0] if dataset == "kaggle" else lstr[-1]
self.npzfile = self.d_path + (
(self.d_file + "_day") if dataset == "kaggle" else self.d_file
)
self.trafile = self.d_path + (
(self.d_file + "_fea") if dataset == "kaggle" else "fea"
)
# check if pre-processed data is available
data_ready = True
if memory_map:
for i in range(days):
reo_data = self.npzfile + "_{0}_reordered.npz".format(i)
if not path.exists(str(reo_data)):
data_ready = False
else:
if not path.exists(str(pro_data)):
data_ready = False
# pre-process data if needed
# WARNNING: when memory mapping is used we get a collection of files
if data_ready:
print("Reading pre-processed data=%s" % (str(pro_data)))
file = str(pro_data)
else:
print("Reading raw data=%s" % (str(raw_path)))
file = data_utils.getCriteoAdData(
raw_path,
out_file,
max_ind_range,
sub_sample_rate,
days,
split,
randomize,
dataset == "kaggle",
memory_map,
dataset_multiprocessing,
)
# get a number of samples per day
total_file = self.d_path + self.d_file + "_day_count.npz"
with np.load(total_file) as data:
total_per_file = data["total_per_file"]
# compute offsets per file
self.offset_per_file = np.array([0] + [x for x in total_per_file])
for i in range(days):
self.offset_per_file[i + 1] += self.offset_per_file[i]
# print(self.offset_per_file)
# setup data
if memory_map:
# setup the training/testing split
self.split = split
if split == 'none' or split == 'train':
self.day = 0
self.max_day_range = days if split == 'none' else days - 1
elif split == 'test' or split == 'val':
self.day = days - 1
num_samples = self.offset_per_file[days] - \
self.offset_per_file[days - 1]
self.test_size = int(np.ceil(num_samples / 2.))
self.val_size = num_samples - self.test_size
else:
sys.exit("ERROR: dataset split is neither none, nor train or test.")
'''
# text
print("text")
for i in range(days):
fi = self.npzfile + "_{0}".format(i)
with open(fi) as data:
ttt = 0; nnn = 0
for _j, line in enumerate(data):
ttt +=1
if np.int32(line[0]) > 0:
nnn +=1
print("day=" + str(i) + " total=" + str(ttt) + " non-zeros="
+ str(nnn) + " ratio=" +str((nnn * 100.) / ttt) + "%")
# processed
print("processed")
for i in range(days):
fi = self.npzfile + "_{0}_processed.npz".format(i)
with np.load(fi) as data:
yyy = data["y"]
ttt = len(yyy)
nnn = np.count_nonzero(yyy)
print("day=" + str(i) + " total=" + str(ttt) + " non-zeros="
+ str(nnn) + " ratio=" +str((nnn * 100.) / ttt) + "%")
# reordered
print("reordered")
for i in range(days):
fi = self.npzfile + "_{0}_reordered.npz".format(i)
with np.load(fi) as data:
yyy = data["y"]
ttt = len(yyy)
nnn = np.count_nonzero(yyy)
print("day=" + str(i) + " total=" + str(ttt) + " non-zeros="
+ str(nnn) + " ratio=" +str((nnn * 100.) / ttt) + "%")
'''
# load unique counts
with np.load(self.d_path + self.d_file + "_fea_count.npz") as data:
self.counts = data["counts"]
self.m_den = den_fea # X_int.shape[1]
self.n_emb = len(self.counts)
print("Sparse features= %d, Dense features= %d" % (self.n_emb, self.m_den))
# Load the test data
# Only a single day is used for testing
if self.split == 'test' or self.split == 'val':
# only a single day is used for testing
fi = self.npzfile + "_{0}_reordered.npz".format(
self.day
)
with np.load(fi) as data:
self.X_int = data["X_int"] # continuous feature
self.X_cat = data["X_cat"] # categorical feature
self.y = data["y"] # target
else:
# load and preprocess data
with np.load(file) as data:
X_int = data["X_int"] # continuous feature
X_cat = data["X_cat"] # categorical feature
y = data["y"] # target
self.counts = data["counts"]
self.m_den = X_int.shape[1] # den_fea
self.n_emb = len(self.counts)
print("Sparse fea = %d, Dense fea = %d" % (self.n_emb, self.m_den))
# create reordering
indices = np.arange(len(y))
if split == "none":
# randomize all data
if randomize == "total":
indices = np.random.permutation(indices)
print("Randomized indices...")
X_int[indices] = X_int
X_cat[indices] = X_cat
y[indices] = y
else:
indices = np.array_split(indices, self.offset_per_file[1:-1])
# randomize train data (per day)
if randomize == "day": # or randomize == "total":
for i in range(len(indices) - 1):
indices[i] = np.random.permutation(indices[i])
print("Randomized indices per day ...")
train_indices = np.concatenate(indices[:-1])
test_indices = indices[-1]
test_indices, val_indices = np.array_split(test_indices, 2)
print("Defined %s indices..." % (split))
# randomize train data (across days)
if randomize == "total":
train_indices = np.random.permutation(train_indices)
print("Randomized indices across days ...")
# create training, validation, and test sets
if split == 'train':
self.X_int = [X_int[i] for i in train_indices]
self.X_cat = [X_cat[i] for i in train_indices]
self.y = [y[i] for i in train_indices]
elif split == 'val':
self.X_int = [X_int[i] for i in val_indices]
self.X_cat = [X_cat[i] for i in val_indices]
self.y = [y[i] for i in val_indices]
elif split == 'test':
self.X_int = [X_int[i] for i in test_indices]
self.X_cat = [X_cat[i] for i in test_indices]
self.y = [y[i] for i in test_indices]
print("Split data according to indices...")
def __getitem__(self, index):
if isinstance(index, slice):
return [
self[idx] for idx in range(
index.start or 0, index.stop or len(self), index.step or 1
)
]
if self.memory_map:
if self.split == 'none' or self.split == 'train':
# check if need to swicth to next day and load data
if index == self.offset_per_file[self.day]:
# print("day_boundary switch", index)
self.day_boundary = self.offset_per_file[self.day]
fi = self.npzfile + "_{0}_reordered.npz".format(
self.day
)
# print('Loading file: ', fi)
with np.load(fi) as data:
self.X_int = data["X_int"] # continuous feature
self.X_cat = data["X_cat"] # categorical feature
self.y = data["y"] # target
self.day = (self.day + 1) % self.max_day_range
i = index - self.day_boundary
elif self.split == 'test' or self.split == 'val':
# only a single day is used for testing
i = index + (0 if self.split == 'test' else self.test_size)
else:
sys.exit("ERROR: dataset split is neither none, nor train or test.")
else:
i = index
if self.max_ind_range > 0:
return self.X_int[i], self.X_cat[i] % self.max_ind_range, self.y[i]
else:
return self.X_int[i], self.X_cat[i], self.y[i]
def _default_preprocess(self, X_int, X_cat, y):
X_int = torch.log(torch.tensor(X_int, dtype=torch.float) + 1)
if self.max_ind_range > 0:
X_cat = torch.tensor(X_cat % self.max_ind_range, dtype=torch.long)
else:
X_cat = torch.tensor(X_cat, dtype=torch.long)
y = torch.tensor(y.astype(np.float32))
return X_int, X_cat, y
def __len__(self):
if self.memory_map:
if self.split == 'none':
return self.offset_per_file[-1]
elif self.split == 'train':
return self.offset_per_file[-2]
elif self.split == 'test':
return self.test_size
elif self.split == 'val':
return self.val_size
else:
sys.exit("ERROR: dataset split is neither none, nor train nor test.")
else:
return len(self.y)
def collate_wrapper_criteo_offset(list_of_tuples):
# where each tuple is (X_int, X_cat, y)
transposed_data = list(zip(*list_of_tuples))
X_int = torch.log(torch.tensor(transposed_data[0], dtype=torch.float) + 1)
X_cat = torch.tensor(transposed_data[1], dtype=torch.long)
T = torch.tensor(transposed_data[2], dtype=torch.float32).view(-1, 1)
batchSize = X_cat.shape[0]
featureCnt = X_cat.shape[1]
lS_i = [X_cat[:, i] for i in range(featureCnt)]
lS_o = [torch.tensor(range(batchSize)) for _ in range(featureCnt)]
return X_int, torch.stack(lS_o), torch.stack(lS_i), T
def ensure_dataset_preprocessed(args, d_path):
_ = CriteoDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"train",
args.raw_data_file,
args.processed_data_file,
args.memory_map,
args.dataset_multiprocessing
)
_ = CriteoDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"test",
args.raw_data_file,
args.processed_data_file,
args.memory_map,
args.dataset_multiprocessing
)
for split in ['train', 'val', 'test']:
print('Running preprocessing for split =', split)
train_files = ['{}_{}_reordered.npz'.format(args.raw_data_file, day)
for
day in range(0, 23)]
test_valid_file = args.raw_data_file + '_23_reordered.npz'
output_file = d_path + '_{}.bin'.format(split)
input_files = train_files if split == 'train' else [test_valid_file]
data_loader_terabyte.numpy_to_binary(input_files=input_files,
output_file_path=output_file,
split=split)
# Conversion from offset to length
def offset_to_length_converter(lS_o, lS_i):
def diff(tensor):
return tensor[1:] - tensor[:-1]
return torch.stack(
[
diff(torch.cat((S_o, torch.tensor(lS_i[ind].shape))).int())
for ind, S_o in enumerate(lS_o)
]
)
def collate_wrapper_criteo_length(list_of_tuples):
# where each tuple is (X_int, X_cat, y)
transposed_data = list(zip(*list_of_tuples))
X_int = torch.log(torch.tensor(transposed_data[0], dtype=torch.float) + 1)
X_cat = torch.tensor(transposed_data[1], dtype=torch.long)
T = torch.tensor(transposed_data[2], dtype=torch.float32).view(-1, 1)
batchSize = X_cat.shape[0]
featureCnt = X_cat.shape[1]
lS_i = torch.stack([X_cat[:, i] for i in range(featureCnt)])
lS_o = torch.stack(
[torch.tensor(range(batchSize)) for _ in range(featureCnt)]
)
lS_l = offset_to_length_converter(lS_o, lS_i)
return X_int, lS_l, lS_i, T
def make_criteo_data_and_loaders(args, offset_to_length_converter=False):
if args.mlperf_logging and args.memory_map and args.data_set == "terabyte":
# more efficient for larger batches
data_directory = path.dirname(args.raw_data_file)
if args.mlperf_bin_loader:
lstr = args.processed_data_file.split("/")
d_path = "/".join(lstr[0:-1]) + "/" + lstr[-1].split(".")[0]
train_file = d_path + "_train.bin"
test_file = d_path + "_test.bin"
# val_file = d_path + "_val.bin"
counts_file = args.raw_data_file + '_fea_count.npz'
if any(not path.exists(p) for p in [train_file,
test_file,
counts_file]):
ensure_dataset_preprocessed(args, d_path)
train_data = data_loader_terabyte.CriteoBinDataset(
data_file=train_file,
counts_file=counts_file,
batch_size=args.mini_batch_size,
max_ind_range=args.max_ind_range
)
mlperf_logger.log_event(key=mlperf_logger.constants.TRAIN_SAMPLES,
value=train_data.num_samples)
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=None,
batch_sampler=None,
shuffle=False,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
sampler=RandomSampler(train_data) if args.mlperf_bin_shuffle else None
)
test_data = data_loader_terabyte.CriteoBinDataset(
data_file=test_file,
counts_file=counts_file,
batch_size=args.test_mini_batch_size,
max_ind_range=args.max_ind_range
)
mlperf_logger.log_event(key=mlperf_logger.constants.EVAL_SAMPLES,
value=test_data.num_samples)
test_loader = torch.utils.data.DataLoader(
test_data,
batch_size=None,
batch_sampler=None,
shuffle=False,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
)
else:
data_filename = args.raw_data_file.split("/")[-1]
train_data = CriteoDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"train",
args.raw_data_file,
args.processed_data_file,
args.memory_map,
args.dataset_multiprocessing
)
test_data = CriteoDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"test",
args.raw_data_file,
args.processed_data_file,
args.memory_map,
args.dataset_multiprocessing
)
train_loader = data_loader_terabyte.DataLoader(
data_directory=data_directory,
data_filename=data_filename,
days=list(range(23)),
batch_size=args.mini_batch_size,
max_ind_range=args.max_ind_range,
split="train"
)
test_loader = data_loader_terabyte.DataLoader(
data_directory=data_directory,
data_filename=data_filename,
days=[23],
batch_size=args.test_mini_batch_size,
max_ind_range=args.max_ind_range,
split="test"
)
else:
train_data = CriteoDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"train",
args.raw_data_file,
args.processed_data_file,
args.memory_map,
args.dataset_multiprocessing,
)
test_data = CriteoDataset(
args.data_set,
args.max_ind_range,
args.data_sub_sample_rate,
args.data_randomize,
"test",
args.raw_data_file,
args.processed_data_file,
args.memory_map,
args.dataset_multiprocessing,
)
collate_wrapper_criteo = collate_wrapper_criteo_offset
if offset_to_length_converter:
collate_wrapper_criteo = collate_wrapper_criteo_length
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=args.mini_batch_size,
shuffle=False,
num_workers=args.num_workers,
collate_fn=collate_wrapper_criteo,
pin_memory=False,
drop_last=False, # True
)
test_loader = torch.utils.data.DataLoader(
test_data,
batch_size=args.test_mini_batch_size,
shuffle=False,
num_workers=args.test_num_workers,
collate_fn=collate_wrapper_criteo,
pin_memory=False,
drop_last=False, # True
)
return train_data, train_loader, test_data, test_loader
# uniform ditribution (input data)
class RandomDataset(Dataset):
def __init__(
self,
m_den,
ln_emb,
data_size,
num_batches,
mini_batch_size,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
num_targets=1,
round_targets=False,
data_generation="random",
trace_file="",
enable_padding=False,
reset_seed_on_access=False,
rand_data_dist="uniform",
rand_data_min=1,
rand_data_max=1,
rand_data_mu=-1,
rand_data_sigma=1,
rand_seed=0
):
# compute batch size
nbatches = int(np.ceil((data_size * 1.0) / mini_batch_size))
if num_batches != 0:
nbatches = num_batches
data_size = nbatches * mini_batch_size
# print("Total number of batches %d" % nbatches)
# save args (recompute data_size if needed)
self.m_den = m_den
self.ln_emb = ln_emb
self.data_size = data_size
self.num_batches = nbatches
self.mini_batch_size = mini_batch_size
self.num_indices_per_lookup = num_indices_per_lookup
self.num_indices_per_lookup_fixed = num_indices_per_lookup_fixed
self.num_targets = num_targets
self.round_targets = round_targets
self.data_generation = data_generation
self.trace_file = trace_file
self.enable_padding = enable_padding
self.reset_seed_on_access = reset_seed_on_access
self.rand_seed = rand_seed
self.rand_data_dist = rand_data_dist
self.rand_data_min = rand_data_min
self.rand_data_max = rand_data_max
self.rand_data_mu = rand_data_mu
self.rand_data_sigma = rand_data_sigma
def reset_numpy_seed(self, numpy_rand_seed):
np.random.seed(numpy_rand_seed)
# torch.manual_seed(numpy_rand_seed)
def __getitem__(self, index):
if isinstance(index, slice):
return [
self[idx] for idx in range(
index.start or 0, index.stop or len(self), index.step or 1
)
]
# WARNING: reset seed on access to first element
# (e.g. if same random samples needed across epochs)
if self.reset_seed_on_access and index == 0:
self.reset_numpy_seed(self.rand_seed)
# number of data points in a batch
n = min(self.mini_batch_size, self.data_size - (index * self.mini_batch_size))
# generate a batch of dense and sparse features
if self.data_generation == "random":
(X, lS_o, lS_i) = generate_dist_input_batch(
self.m_den,
self.ln_emb,
n,
self.num_indices_per_lookup,
self.num_indices_per_lookup_fixed,
rand_data_dist=self.rand_data_dist,
rand_data_min=self.rand_data_min,
rand_data_max=self.rand_data_max,
rand_data_mu=self.rand_data_mu,
rand_data_sigma=self.rand_data_sigma,
)
elif self.data_generation == "synthetic":
(X, lS_o, lS_i) = generate_synthetic_input_batch(
self.m_den,
self.ln_emb,
n,
self.num_indices_per_lookup,
self.num_indices_per_lookup_fixed,
self.trace_file,
self.enable_padding
)
else:
sys.exit(
"ERROR: --data-generation=" + self.data_generation + " is not supported"
)
# generate a batch of target (probability of a click)
T = generate_random_output_batch(n, self.num_targets, self.round_targets)
return (X, lS_o, lS_i, T)
def __len__(self):
# WARNING: note that we produce bacthes of outputs in __getitem__
# therefore we should use num_batches rather than data_size below
return self.num_batches
def collate_wrapper_random_offset(list_of_tuples):
# where each tuple is (X, lS_o, lS_i, T)
(X, lS_o, lS_i, T) = list_of_tuples[0]
return (X,
torch.stack(lS_o),
lS_i,
T)
def collate_wrapper_random_length(list_of_tuples):
# where each tuple is (X, lS_o, lS_i, T)
(X, lS_o, lS_i, T) = list_of_tuples[0]
return (X,
offset_to_length_converter(torch.stack(lS_o), lS_i),
lS_i,
T)
def make_random_data_and_loader(args, ln_emb, m_den,
offset_to_length_converter=False,
):
train_data = RandomDataset(
m_den,
ln_emb,
args.data_size,
args.num_batches,
args.mini_batch_size,
args.num_indices_per_lookup,
args.num_indices_per_lookup_fixed,
1, # num_targets
args.round_targets,
args.data_generation,
args.data_trace_file,
args.data_trace_enable_padding,
reset_seed_on_access=True,
rand_data_dist=args.rand_data_dist,
rand_data_min=args.rand_data_min,
rand_data_max=args.rand_data_max,
rand_data_mu=args.rand_data_mu,
rand_data_sigma=args.rand_data_sigma,
rand_seed=args.numpy_rand_seed
) # WARNING: generates a batch of lookups at once
test_data = RandomDataset(
m_den,
ln_emb,
args.data_size,
args.num_batches,
args.mini_batch_size,
args.num_indices_per_lookup,
args.num_indices_per_lookup_fixed,
1, # num_targets
args.round_targets,
args.data_generation,
args.data_trace_file,
args.data_trace_enable_padding,
reset_seed_on_access=True,
rand_data_dist=args.rand_data_dist,
rand_data_min=args.rand_data_min,
rand_data_max=args.rand_data_max,
rand_data_mu=args.rand_data_mu,
rand_data_sigma=args.rand_data_sigma,
rand_seed=args.numpy_rand_seed
)
collate_wrapper_random = collate_wrapper_random_offset
if offset_to_length_converter:
collate_wrapper_random = collate_wrapper_random_length
train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=1,
shuffle=False,
num_workers=args.num_workers,
collate_fn=collate_wrapper_random,
pin_memory=False,
drop_last=False, # True
)
test_loader = torch.utils.data.DataLoader(
test_data,
batch_size=1,
shuffle=False,
num_workers=args.num_workers,
collate_fn=collate_wrapper_random,
pin_memory=False,
drop_last=False, # True
)
return train_data, train_loader, test_data, test_loader
def generate_random_data(
m_den,
ln_emb,
data_size,
num_batches,
mini_batch_size,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
num_targets=1,
round_targets=False,
data_generation="random",
trace_file="",
enable_padding=False,
length=False, # length for caffe2 version (except dlrm_s_caffe2)
):
nbatches = int(np.ceil((data_size * 1.0) / mini_batch_size))
if num_batches != 0:
nbatches = num_batches
data_size = nbatches * mini_batch_size
# print("Total number of batches %d" % nbatches)
# inputs
lT = []
lX = []
lS_offsets = []
lS_indices = []
for j in range(0, nbatches):
# number of data points in a batch
n = min(mini_batch_size, data_size - (j * mini_batch_size))
# generate a batch of dense and sparse features
if data_generation == "random":
(Xt, lS_emb_offsets, lS_emb_indices) = generate_uniform_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
length,
)
elif data_generation == "synthetic":
(Xt, lS_emb_offsets, lS_emb_indices) = generate_synthetic_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
trace_file,
enable_padding
)
else:
sys.exit(
"ERROR: --data-generation=" + data_generation + " is not supported"
)
# dense feature
lX.append(Xt)
# sparse feature (sparse indices)
lS_offsets.append(lS_emb_offsets)
lS_indices.append(lS_emb_indices)
# generate a batch of target (probability of a click)
P = generate_random_output_batch(n, num_targets, round_targets)
lT.append(P)
return (nbatches, lX, lS_offsets, lS_indices, lT)
def generate_random_output_batch(n, num_targets, round_targets=False):
# target (probability of a click)
if round_targets:
P = np.round(ra.rand(n, num_targets).astype(np.float32)).astype(np.float32)
else:
P = ra.rand(n, num_targets).astype(np.float32)
return torch.tensor(P)
# uniform ditribution (input data)
def generate_uniform_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
length,
):
# dense feature
Xt = torch.tensor(ra.rand(n, m_den).astype(np.float32))
# sparse feature (sparse indices)
lS_emb_offsets = []
lS_emb_indices = []
# for each embedding generate a list of n lookups,
# where each lookup is composed of multiple sparse indices
for size in ln_emb:
lS_batch_offsets = []
lS_batch_indices = []
offset = 0
for _ in range(n):
# num of sparse indices to be used per embedding (between
if num_indices_per_lookup_fixed:
sparse_group_size = np.int64(num_indices_per_lookup)
else:
# random between [1,num_indices_per_lookup])
r = ra.random(1)
sparse_group_size = np.int64(
np.round(max([1.0], r * min(size, num_indices_per_lookup)))
)
# sparse indices to be used per embedding
r = ra.random(sparse_group_size)
sparse_group = np.unique(np.round(r * (size - 1)).astype(np.int64))
# reset sparse_group_size in case some index duplicates were removed
sparse_group_size = np.int32(sparse_group.size)
# store lengths and indices
if length: # for caffe2 version
lS_batch_offsets += [sparse_group_size]
else:
lS_batch_offsets += [offset]
lS_batch_indices += sparse_group.tolist()
# update offset for next iteration
offset += sparse_group_size
lS_emb_offsets.append(torch.tensor(lS_batch_offsets))
lS_emb_indices.append(torch.tensor(lS_batch_indices))
return (Xt, lS_emb_offsets, lS_emb_indices)
# random data from uniform or gaussian ditribution (input data)
def generate_dist_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
rand_data_dist,
rand_data_min,
rand_data_max,
rand_data_mu,
rand_data_sigma,
):
# dense feature
Xt = torch.tensor(ra.rand(n, m_den).astype(np.float32))
# sparse feature (sparse indices)
lS_emb_offsets = []
lS_emb_indices = []
# for each embedding generate a list of n lookups,
# where each lookup is composed of multiple sparse indices
for size in ln_emb:
lS_batch_offsets = []
lS_batch_indices = []
offset = 0
for _ in range(n):
# num of sparse indices to be used per embedding (between
if num_indices_per_lookup_fixed:
sparse_group_size = np.int64(num_indices_per_lookup)
else:
# random between [1,num_indices_per_lookup])
r = ra.random(1)
sparse_group_size = np.int64(
np.round(max([1.0], r * min(size, num_indices_per_lookup)))
)
# sparse indices to be used per embedding
if rand_data_dist == "gaussian":
if rand_data_mu == -1:
rand_data_mu = (rand_data_max + rand_data_min) / 2.0
r = ra.normal(rand_data_mu, rand_data_sigma, sparse_group_size)
sparse_group = np.clip(r, rand_data_min, rand_data_max)
sparse_group = np.unique(sparse_group).astype(np.int64)
elif rand_data_dist == "uniform":
r = ra.random(sparse_group_size)
sparse_group = np.unique(np.round(r * (size - 1)).astype(np.int64))
else:
raise(rand_data_dist, "distribution is not supported. \
please select uniform or gaussian")
# reset sparse_group_size in case some index duplicates were removed
sparse_group_size = np.int64(sparse_group.size)
# store lengths and indices
lS_batch_offsets += [offset]
lS_batch_indices += sparse_group.tolist()
# update offset for next iteration
offset += sparse_group_size
lS_emb_offsets.append(torch.tensor(lS_batch_offsets))
lS_emb_indices.append(torch.tensor(lS_batch_indices))
return (Xt, lS_emb_offsets, lS_emb_indices)
# synthetic distribution (input data)
def generate_synthetic_input_batch(
m_den,
ln_emb,
n,
num_indices_per_lookup,
num_indices_per_lookup_fixed,
trace_file,
enable_padding=False,
):
# dense feature
Xt = torch.tensor(ra.rand(n, m_den).astype(np.float32))
# sparse feature (sparse indices)
lS_emb_offsets = []
lS_emb_indices = []
# for each embedding generate a list of n lookups,
# where each lookup is composed of multiple sparse indices
for i, size in enumerate(ln_emb):
lS_batch_offsets = []
lS_batch_indices = []
offset = 0
for _ in range(n):
# num of sparse indices to be used per embedding (between
if num_indices_per_lookup_fixed:
sparse_group_size = np.int64(num_indices_per_lookup)
else:
# random between [1,num_indices_per_lookup])
r = ra.random(1)
sparse_group_size = np.int64(
max(1, np.round(r * min(size, num_indices_per_lookup))[0])
)
# sparse indices to be used per embedding
file_path = trace_file
line_accesses, list_sd, cumm_sd = read_dist_from_file(
file_path.replace("j", str(i))
)
# debug prints
# print("input")
# print(line_accesses); print(list_sd); print(cumm_sd);
# print(sparse_group_size)
# approach 1: rand
# r = trace_generate_rand(
# line_accesses, list_sd, cumm_sd, sparse_group_size, enable_padding
# )
# approach 2: lru
r = trace_generate_lru(
line_accesses, list_sd, cumm_sd, sparse_group_size, enable_padding
)
# WARNING: if the distribution in the file is not consistent
# with embedding table dimensions, below mod guards against out
# of range access
sparse_group = np.unique(r).astype(np.int64)
minsg = np.min(sparse_group)
maxsg = np.max(sparse_group)
if (minsg < 0) or (size <= maxsg):
print(
"WARNING: distribution is inconsistent with embedding "
+ "table size (using mod to recover and continue)"
)
sparse_group = np.mod(sparse_group, size).astype(np.int64)
# sparse_group = np.unique(np.array(np.mod(r, size-1)).astype(np.int64))
# reset sparse_group_size in case some index duplicates were removed
sparse_group_size = np.int64(sparse_group.size)
# store lengths and indices
lS_batch_offsets += [offset]
lS_batch_indices += sparse_group.tolist()
# update offset for next iteration
offset += sparse_group_size
lS_emb_offsets.append(torch.tensor(lS_batch_offsets))
lS_emb_indices.append(torch.tensor(lS_batch_indices))
return (Xt, lS_emb_offsets, lS_emb_indices)
def generate_stack_distance(cumm_val, cumm_dist, max_i, i, enable_padding=False):
u = ra.rand(1)
if i < max_i:
# only generate stack distances up to the number of new references seen so far
j = bisect.bisect(cumm_val, i) - 1
fi = cumm_dist[j]
u *= fi # shrink distribution support to exclude last values
elif enable_padding:
# WARNING: disable generation of new references (once all have been seen)
fi = cumm_dist[0]
u = (1.0 - fi) * u + fi # remap distribution support to exclude first value
for (j, f) in enumerate(cumm_dist):
if u <= f:
return cumm_val[j]
# WARNING: global define, must be consistent across all synthetic functions
cache_line_size = 1
def trace_generate_lru(
line_accesses, list_sd, cumm_sd, out_trace_len, enable_padding=False
):
max_sd = list_sd[-1]
l = len(line_accesses)
i = 0
ztrace = deque()
for _ in range(out_trace_len):
sd = generate_stack_distance(list_sd, cumm_sd, max_sd, i, enable_padding)
mem_ref_within_line = 0 # floor(ra.rand(1)*cache_line_size) #0
# generate memory reference
if sd == 0: # new reference #
line_ref = line_accesses[0]
del line_accesses[0]
line_accesses.append(line_ref)
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
i += 1
else: # existing reference #
line_ref = line_accesses[l - sd]
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
del line_accesses[l - sd]
line_accesses.append(line_ref)
# save generated memory reference
ztrace.append(mem_ref)
return ztrace
def trace_generate_rand(
line_accesses, list_sd, cumm_sd, out_trace_len, enable_padding=False
):
max_sd = list_sd[-1]
l = len(line_accesses) # !!!Unique,
i = 0
ztrace = []
for _ in range(out_trace_len):
sd = generate_stack_distance(list_sd, cumm_sd, max_sd, i, enable_padding)
mem_ref_within_line = 0 # floor(ra.rand(1)*cache_line_size) #0
# generate memory reference
if sd == 0: # new reference #
line_ref = line_accesses.pop(0)
line_accesses.append(line_ref)
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
i += 1
else: # existing reference #
line_ref = line_accesses[l - sd]
mem_ref = np.uint64(line_ref * cache_line_size + mem_ref_within_line)
ztrace.append(mem_ref)
return ztrace
def trace_profile(trace, enable_padding=False):
# number of elements in the array (assuming 1D)
# n = trace.size
rstack = deque() # S
stack_distances = deque() # SDS
line_accesses = deque() # L
for x in trace:
r = np.uint64(x / cache_line_size)
l = len(rstack)
try: # found #
i = rstack.index(r)
# WARNING: I believe below is the correct depth in terms of meaning of the
# algorithm, but that is not what seems to be in the paper alg.
# -1 can be subtracted if we defined the distance between
# consecutive accesses (e.g. r, r) as 0 rather than 1.
sd = l - i # - 1
# push r to the end of stack_distances
stack_distances.appendleft(sd)
# remove r from its position and insert to the top of stack
del rstack[i] # rstack.remove(r)
rstack.append(r)
except ValueError: # not found #
sd = 0 # -1
# push r to the end of stack_distances/line_accesses
stack_distances.appendleft(sd)
line_accesses.appendleft(r)
# push r to the top of stack
rstack.append(r)
if enable_padding:
# WARNING: notice that as the ratio between the number of samples (l)
# and cardinality (c) of a sample increases the probability of
# generating a sample gets smaller and smaller because there are
# few new samples compared to repeated samples. This means that for a
# long trace with relatively small cardinality it will take longer to
# generate all new samples and therefore obtain full distribution support
# and hence it takes longer for distribution to resemble the original.
# Therefore, we may pad the number of new samples to be on par with
# average number of samples l/c artificially.
l = len(stack_distances)
c = max(stack_distances)
padding = int(np.ceil(l / c))
stack_distances = stack_distances + [0] * padding
return (rstack, stack_distances, line_accesses)
# auxiliary read/write routines
def read_trace_from_file(file_path):
try:
with open(file_path) as f:
if args.trace_file_binary_type:
array = np.fromfile(f, dtype=np.uint64)
trace = array.astype(np.uint64).tolist()
else:
line = f.readline()
trace = list(map(lambda x: np.uint64(x), line.split(", ")))
return trace
except Exception:
print(f"ERROR: trace file '{file_path}' is not available.")
def write_trace_to_file(file_path, trace):
try:
if args.trace_file_binary_type:
with open(file_path, "wb+") as f:
np.array(trace).astype(np.uint64).tofile(f)
else:
with open(file_path, "w+") as f:
s = str(list(trace))
f.write(s[1 : len(s) - 1])
except Exception:
print("ERROR: no output trace file has been provided")
def read_dist_from_file(file_path):
try:
with open(file_path, "r") as f:
lines = f.read().splitlines()
except Exception:
print("{file_path} Wrong file or file path")
# read unique accesses
unique_accesses = [int(el) for el in lines[0].split(", ")]
# read cumulative distribution (elements are passed as two separate lists)
list_sd = [int(el) for el in lines[1].split(", ")]
cumm_sd = [float(el) for el in lines[2].split(", ")]
return unique_accesses, list_sd, cumm_sd
def write_dist_to_file(file_path, unique_accesses, list_sd, cumm_sd):
try:
with open(file_path, "w") as f:
# unique_acesses
s = str(list(unique_accesses))
f.write(s[1 : len(s) - 1] + "\n")
# list_sd
s = str(list_sd)
f.write(s[1 : len(s) - 1] + "\n")
# cumm_sd
s = str(list(cumm_sd))
f.write(s[1 : len(s) - 1] + "\n")
except Exception:
print("Wrong file or file path")
if __name__ == "__main__":
import operator
import argparse
### parse arguments ###
parser = argparse.ArgumentParser(description="Generate Synthetic Distributions")
parser.add_argument("--trace-file", type=str, default="./input/trace.log")
parser.add_argument("--trace-file-binary-type", type=bool, default=False)
parser.add_argument("--trace-enable-padding", type=bool, default=False)
parser.add_argument("--dist-file", type=str, default="./input/dist.log")
parser.add_argument(
"--synthetic-file", type=str, default="./input/trace_synthetic.log"
)
parser.add_argument("--numpy-rand-seed", type=int, default=123)
parser.add_argument("--print-precision", type=int, default=5)
args = parser.parse_args()
### some basic setup ###
np.random.seed(args.numpy_rand_seed)
np.set_printoptions(precision=args.print_precision)
### read trace ###
trace = read_trace_from_file(args.trace_file)
# print(trace)
### profile trace ###
(_, stack_distances, line_accesses) = trace_profile(
trace, args.trace_enable_padding
)
stack_distances.reverse()
line_accesses.reverse()
# print(line_accesses)
# print(stack_distances)
### compute probability distribution ###
# count items
l = len(stack_distances)
dc = sorted(
collections.Counter(stack_distances).items(), key=operator.itemgetter(0)
)
# create a distribution
list_sd = list(map(lambda tuple_x_k: tuple_x_k[0], dc)) # x = tuple_x_k[0]
dist_sd = list(
map(lambda tuple_x_k: tuple_x_k[1] / float(l), dc)
) # k = tuple_x_k[1]
cumm_sd = deque() # np.cumsum(dc).tolist() #prefixsum
for i, (_, k) in enumerate(dc):
if i == 0:
cumm_sd.append(k / float(l))
else:
# add the 2nd element of the i-th tuple in the dist_sd list
cumm_sd.append(cumm_sd[i - 1] + (k / float(l)))
### write stack_distance and line_accesses to a file ###
write_dist_to_file(args.dist_file, line_accesses, list_sd, cumm_sd)
### generate corresponding synthetic ###
# line_accesses, list_sd, cumm_sd = read_dist_from_file(args.dist_file)
synthetic_trace = trace_generate_lru(
line_accesses, list_sd, cumm_sd, len(trace), args.trace_enable_padding
)
# synthetic_trace = trace_generate_rand(
# line_accesses, list_sd, cumm_sd, len(trace), args.trace_enable_padding
# )
write_trace_to_file(args.synthetic_file, synthetic_trace)
PyTorch/Recommendation/dlrm/dlrm_s_caffe2.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: an implementation of a deep learning recommendation model (DLRM)
# The model input consists of dense and sparse features. The former is a vector
# of floating point values. The latter is a list of sparse indices into
# embedding tables, which consist of vectors of floating point values.
# The selected vectors are passed to mlp networks denoted by triangles,
# in some cases the vectors are interacted through operators (Ops).
#
# output:
# vector of values
# model: |
# /\
# /__\
# |
# _____________________> Op <___________________
# / | \
# /\ /\ /\
# /__\ /__\ ... /__\
# | | |
# | Op Op
# | ____/__\_____ ____/__\____
# | |_Emb_|____|__| ... |_Emb_|__|___|
# input:
# [ dense features ] [sparse indices] , ..., [sparse indices]
#
# More precise definition of model layers:
# 1) fully connected layers of an mlp
# z = f(y)
# y = Wx + b
#
# 2) embedding lookup (for a list of sparse indices p=[p1,...,pk])
# z = Op(e1,...,ek)
# obtain vectors e1=E[:,p1], ..., ek=E[:,pk]
#
# 3) Operator Op can be one of the following
# Sum(e1,...,ek) = e1 + ... + ek
# Dot(e1,...,ek) = [e1'e1, ..., e1'ek, ..., ek'e1, ..., ek'ek]
# Cat(e1,...,ek) = [e1', ..., ek']'
# where ' denotes transpose operation
#
# References:
# [1] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang,
# Narayanan Sundaram, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu,
# Alisson G. Azzolini, Dmytro Dzhulgakov, Andrey Mallevich, Ilia Cherniavskii,
# Yinghai Lu, Raghuraman Krishnamoorthi, Ansha Yu, Volodymyr Kondratenko,
# Stephanie Pereira, Xianjie Chen, Wenlin Chen, Vijay Rao, Bill Jia, Liang Xiong,
# Misha Smelyanskiy, "Deep Learning Recommendation Model for Personalization and
# Recommendation Systems", CoRR, arXiv:1906.00091, 2019
from __future__ import absolute_import, division, print_function, unicode_literals
import functools
# others
import operator
import time
import copy
# data generation
import dlrm_data_pytorch as dp
# numpy
import numpy as np
import sklearn.metrics
# onnx
# The onnx import causes deprecation warnings every time workers
# are spawned during testing. So, we filter out those warnings.
import warnings
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=DeprecationWarning)
try:
import onnx
import caffe2.python.onnx.frontend
except ImportError as error:
print('Unable to import onnx or caffe2.python.onnx.frontend ', error)
# from caffe2.python import data_parallel_model
# caffe2
from caffe2.proto import caffe2_pb2
from caffe2.python import brew, core, dyndep, model_helper, net_drawer, workspace
"""
# auxiliary routine used to split input on the mini-bacth dimension
def where_to_split(mini_batch_size, ndevices, _add_leftover=False):
n = (mini_batch_size + ndevices - 1) // ndevices # ceiling
l = mini_batch_size - n * (ndevices - 1) # leftover
s = [n] * (ndevices - 1)
if _add_leftover:
ls += [l if l > 0 else n]
return ls
"""
### define dlrm in Caffe2 ###
class DLRM_Net(object):
def FeedBlobWrapper(self, tag, val, add_prefix=True, split=False, device_id=-1):
if self.ndevices > 1 and add_prefix:
if split:
# split across devices
mini_batch_size = val.shape[0]
# approach 1: np and caffe2 operators assume the mini-batch size is
# divisible exactly by the number of available devices
if mini_batch_size % self.ndevices != 0:
sys.exit("ERROR: caffe2 net assumes that the mini_batch_size "
+ str(mini_batch_size)
+ " is evenly divisible by the number of available devices"
+ str(self.ndevices))
vals = np.split(val, self.ndevices, axis=0)
"""
# approach 2: np and caffe2 operators do not assume exact divisibility
if args.mini_batch_size != mini_batch_size:
sys.exit("ERROR: caffe2 net was prepared for mini-batch size "
+ str(args.mini_batch_size)
+ " which is different from current mini-batch size "
+ str(mini_batch_size) + " being passed to it. "
+ "This is common for the last mini-batch, when "
+ "mini-batch size does not evenly divided the number of "
+ "elements in the data set.")
ls = where_to_split(mini_batch_size, self.ndevices)
vals = np.split(val, ls, axis=0)
"""
# feed to multiple devices
for d in range(self.ndevices):
tag_on_device = "gpu_" + str(d) + "/" + tag
_d = core.DeviceOption(workspace.GpuDeviceType, d)
workspace.FeedBlob(tag_on_device, vals[d], device_option=_d)
else:
# feed to multiple devices
for d in range(self.ndevices):
tag_on_device = "gpu_" + str(d) + "/" + tag
_d = core.DeviceOption(workspace.GpuDeviceType, d)
workspace.FeedBlob(tag_on_device, val, device_option=_d)
else:
# feed to a single device (named or not)
if device_id >= 0:
_d = core.DeviceOption(workspace.GpuDeviceType, device_id)
workspace.FeedBlob(tag, val, device_option=_d)
else:
workspace.FeedBlob(tag, val)
def FetchBlobWrapper(self, tag, add_prefix=True, reduce_across=None, device_id=-1):
if self.ndevices > 1 and add_prefix:
# fetch from multiple devices
vals = []
for d in range(self.ndevices):
if tag.__class__ == list:
tag_on_device = tag[d]
else:
tag_on_device = "gpu_" + str(0) + "/" + tag
val = workspace.FetchBlob(tag_on_device)
vals.append(val)
# reduce across devices
if reduce_across == "add":
return functools.reduce(operator.add, vals)
elif reduce_across == "concat":
return np.concatenate(vals)
else:
return vals
else:
# fetch from a single device (named or not)
if device_id >= 0:
tag_on_device = "gpu_" + str(device_id) + "/" + tag
return workspace.FetchBlob(tag_on_device)
else:
return workspace.FetchBlob(tag)
def AddLayerWrapper(self, layer, inp_blobs, out_blobs,
add_prefix=True, reset_grad=False, **kwargs):
# auxiliary routine to adjust tags
def adjust_tag(blobs, on_device):
if blobs.__class__ == str:
_blobs = on_device + blobs
elif blobs.__class__ == list:
_blobs = list(map(lambda tag: on_device + tag, blobs))
else: # blobs.__class__ == model_helper.ModelHelper or something else
_blobs = blobs
return _blobs
if self.ndevices > 1 and add_prefix:
# add layer on multiple devices
ll = []
for d in range(self.ndevices):
# add prefix on_device
on_device = "gpu_" + str(d) + "/"
_inp_blobs = adjust_tag(inp_blobs, on_device)
_out_blobs = adjust_tag(out_blobs, on_device)
# WARNING: reset_grad option was exlusively designed for WeightedSum
# with inp_blobs=[w, tag_one, "", lr], where "" will be replaced
if reset_grad:
w_grad = self.gradientMap[_inp_blobs[0]]
_inp_blobs[2] = w_grad
# add layer to the model
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, d)):
if kwargs:
new_layer = layer(_inp_blobs, _out_blobs, **kwargs)
else:
new_layer = layer(_inp_blobs, _out_blobs)
ll.append(new_layer)
return ll
else:
# add layer on a single device
# WARNING: reset_grad option was exlusively designed for WeightedSum
# with inp_blobs=[w, tag_one, "", lr], where "" will be replaced
if reset_grad:
w_grad = self.gradientMap[inp_blobs[0]]
inp_blobs[2] = w_grad
# add layer to the model
if kwargs:
new_layer = layer(inp_blobs, out_blobs, **kwargs)
else:
new_layer = layer(inp_blobs, out_blobs)
return new_layer
def create_mlp(self, ln, sigmoid_layer, model, tag):
(tag_layer, tag_in, tag_out) = tag
# build MLP layer by layer
layers = []
weights = []
for i in range(1, ln.size):
n = ln[i - 1]
m = ln[i]
# create tags
tag_fc_w = tag_layer + ":::" + "fc" + str(i) + "_w"
tag_fc_b = tag_layer + ":::" + "fc" + str(i) + "_b"
tag_fc_y = tag_layer + ":::" + "fc" + str(i) + "_y"
tag_fc_z = tag_layer + ":::" + "fc" + str(i) + "_z"
if i == ln.size - 1:
tag_fc_z = tag_out
weights.append(tag_fc_w)
weights.append(tag_fc_b)
# initialize the weights
# approach 1: custom Xavier input, output or two-sided fill
mean = 0.0 # std_dev = np.sqrt(variance)
std_dev = np.sqrt(2 / (m + n)) # np.sqrt(1 / m) # np.sqrt(1 / n)
W = np.random.normal(mean, std_dev, size=(m, n)).astype(np.float32)
std_dev = np.sqrt(1 / m) # np.sqrt(2 / (m + 1))
b = np.random.normal(mean, std_dev, size=m).astype(np.float32)
self.FeedBlobWrapper(tag_fc_w, W)
self.FeedBlobWrapper(tag_fc_b, b)
# approach 2: caffe2 xavier
# W = self.AddLayerWrapper(
# model.param_init_net.XavierFill,
# [],
# tag_fc_w,
# shape=[m, n]
# )
# b = self.AddLayerWrapper(
# model.param_init_net.ConstantFill,
# [],
# tag_fc_b,
# shape=[m]
# )
# initialize the MLP's momentum for the Adagrad optimizer
if self.emb_optimizer in ["adagrad", "rwsadagrad"]:
# momentum of the weights
self.FeedBlobWrapper(
"momentum_mlp_{}_{}".format(tag_layer, 2 * i - 1),
np.full((m, n), 0, dtype=np.float32)
)
# momentum of the biases
self.FeedBlobWrapper(
"momentum_mlp_{}_{}".format(tag_layer, 2 * i),
np.full((m), 0, dtype=np.float32)
)
# save the blob shapes for latter (only needed if onnx is requested)
if self.save_onnx:
self.onnx_tsd[tag_fc_w] = (onnx.TensorProto.FLOAT, W.shape)
self.onnx_tsd[tag_fc_b] = (onnx.TensorProto.FLOAT, b.shape)
# approach 1: construct fully connected operator using model.net
fc = self.AddLayerWrapper(
model.net.FC, [tag_in, tag_fc_w, tag_fc_b], tag_fc_y
)
# approach 2: construct fully connected operator using brew
# https://github.com/caffe2/tutorials/blob/master/MNIST.ipynb
# fc = brew.fc(model, layer, tag_fc_w, dim_in=m, dim_out=n)
layers.append(fc)
if i == sigmoid_layer:
# approach 1: construct sigmoid operator using model.net
layer = self.AddLayerWrapper(model.net.Sigmoid, tag_fc_y, tag_fc_z)
# approach 2: using brew (which currently does not support sigmoid)
# tag_sigm = tag_layer + ":::" + "sigmoid" + str(i)
# layer = brew.sigmoid(model,fc,tag_sigmoid)
else:
# approach 1: construct relu operator using model.net
layer = self.AddLayerWrapper(model.net.Relu, tag_fc_y, tag_fc_z)
# approach 2: using brew
# tag_relu = tag_layer + ":::" + "relu" + str(i)
# layer = brew.relu(model,fc,tag_relu)
tag_in = tag_fc_z
layers.append(layer)
# WARNING: the dependency between layers is implicit in the tags,
# so only the last layer is added to the layers list. It will
# later be used for interactions.
return layers, weights
def create_emb(self, m, ln, model, tag):
(tag_layer, tag_in, tag_out) = tag
emb_l = []
weights_l = []
vw_l = []
for i in range(0, ln.size):
n = ln[i]
# select device
if self.ndevices > 1:
d = i % self.ndevices
else:
d = -1
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
len_s = on_device + tag_layer + ":::" + "sls" + str(i) + "_l"
ind_s = on_device + tag_layer + ":::" + "sls" + str(i) + "_i"
tbl_s = on_device + tag_layer + ":::" + "sls" + str(i) + "_w"
sum_s = on_device + tag_layer + ":::" + "sls" + str(i) + "_z"
weights_l.append(tbl_s)
# initialize the weights
# approach 1a: custom
W = np.random.uniform(low=-np.sqrt(1 / n),
high=np.sqrt(1 / n),
size=(n, m)).astype(np.float32)
# approach 1b: numpy rand
# W = ra.rand(n, m).astype(np.float32)
self.FeedBlobWrapper(tbl_s, W, False, device_id=d)
# approach 2: caffe2 xavier
# with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, d)):
# W = model.param_init_net.XavierFill([], tbl_s, shape=[n, m])
# save the blob shapes for latter (only needed if onnx is requested)
# initialize the embedding's momentum for the Adagrad optimizer
if self.emb_optimizer == "adagrad":
self.FeedBlobWrapper("momentum_emb_{}".format(i),
np.full((n, m), 0), add_prefix=False, device_id=d)
elif self.emb_optimizer == "rwsadagrad":
self.FeedBlobWrapper("momentum_emb_{}".format(i),
np.full((n), 0), add_prefix=False, device_id=d)
if self.save_onnx:
self.onnx_tsd[tbl_s] = (onnx.TensorProto.FLOAT, W.shape)
# create operator
if self.weighted_pooling is not None:
vw_s = on_device + tag_layer + ":::" + "sls" + str(i) + "_v"
psw_s = on_device + tag_layer + ":::" + "sls" + str(i) + "_s"
VW = np.ones(n).astype(np.float32)
self.FeedBlobWrapper(vw_s, VW, False, device_id=d)
if self.weighted_pooling == "learned":
vw_l.append(vw_s)
grad_on_weights = True
else:
grad_on_weights = False
if self.save_onnx:
self.onnx_tsd[vw_s] = (onnx.TensorProto.FLOAT, VW.shape)
if self.ndevices <= 1:
PSW = model.net.Gather([vw_s, ind_s], [psw_s])
EE = model.net.SparseLengthsWeightedSum(
[tbl_s, PSW, ind_s, len_s], [sum_s],
grad_on_weights=grad_on_weights
)
else:
with core.DeviceScope(
core.DeviceOption(workspace.GpuDeviceType, d)
):
PSW = model.net.Gather([vw_s, ind_s], [psw_s])
EE = model.net.SparseLengthsWeightedSum(
[tbl_s, PSW, ind_s, len_s], [sum_s],
grad_on_weights=grad_on_weights
)
else:
if self.ndevices <= 1:
EE = model.net.SparseLengthsSum(
[tbl_s, ind_s, len_s], [sum_s]
)
else:
with core.DeviceScope(
core.DeviceOption(workspace.GpuDeviceType, d)
):
EE = model.net.SparseLengthsSum(
[tbl_s, ind_s, len_s], [sum_s]
)
emb_l.append(EE)
return emb_l, weights_l, vw_l
def create_interactions(self, x, ly, model, tag):
(tag_dense_in, tag_sparse_in, tag_int_out) = tag
if self.arch_interaction_op == "dot":
# concatenate dense and sparse features
tag_int_out_info = tag_int_out + "_info"
T, T_info = model.net.Concat(
x + ly,
[tag_int_out + "_cat_axis0", tag_int_out_info + "_cat_axis0"],
axis=1,
add_axis=1,
)
# perform a dot product
Z = model.net.BatchMatMul([T, T], tag_int_out + "_matmul", trans_b=1)
# append dense feature with the interactions (into a row vector)
# approach 1: all
# Zflat = model.net.Flatten(Z, tag_int_out + "_flatten", axis=1)
# approach 2: unique
Zflat_all = model.net.Flatten(Z, tag_int_out + "_flatten_all", axis=1)
Zflat = model.net.BatchGather(
[Zflat_all, tag_int_out + "_tril_indices"],
tag_int_out + "_flatten"
)
R, R_info = model.net.Concat(
x + [Zflat], [tag_int_out, tag_int_out_info], axis=1
)
elif self.arch_interaction_op == "cat":
# concatenation features (into a row vector)
tag_int_out_info = tag_int_out + "_info"
R, R_info = model.net.Concat(
x + ly, [tag_int_out, tag_int_out_info], axis=1
)
else:
sys.exit("ERROR: --arch-interaction-op="
+ self.arch_interaction_op + " is not supported")
return R
def create_sequential_forward_ops(self):
# embeddings
tag = (self.temb, self.tsin, self.tsout)
self.emb_l, self.emb_w, self.emb_vw = self.create_emb(
self.m_spa, self.ln_emb, self.model, tag
)
# bottom mlp
tag = (self.tbot, self.tdin, self.tdout)
self.bot_l, self.bot_w = self.create_mlp(self.ln_bot, self.sigmoid_bot,
self.model, tag)
# interactions
tag = (self.tdout, self.tsout, self.tint)
Z = self.create_interactions([self.bot_l[-1]], self.emb_l, self.model, tag)
# top mlp
tag = (self.ttop, Z, self.tout)
self.top_l, self.top_w = self.create_mlp(self.ln_top, self.sigmoid_top,
self.model, tag)
# debug prints
# print(self.emb_l)
# print(self.bot_l)
# print(self.top_l)
# setup the last output variable
self.last_output = self.top_l[-1]
def create_parallel_forward_ops(self):
# distribute embeddings (model parallelism)
tag = (self.temb, self.tsin, self.tsout)
self.emb_l, self.emb_w, self.emb_vw = self.create_emb(
self.m_spa, self.ln_emb, self.model, tag
)
# replicate mlp (data parallelism)
tag = (self.tbot, self.tdin, self.tdout)
self.bot_l, self.bot_w = self.create_mlp(self.ln_bot, self.sigmoid_bot,
self.model, tag)
# add communication (butterfly shuffle)
t_list = []
for i, emb_output in enumerate(self.emb_l):
# split input
src_d = i % self.ndevices
lo = [emb_output + "_split_" + str(d) for d in range(self.ndevices)]
# approach 1: np and caffe2 operators assume the mini-batch size is
# divisible exactly by the number of available devices
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, src_d)):
self.model.net.Split(emb_output, lo, axis=0)
"""
# approach 2: np and caffe2 operators do not assume exact divisibility
ls = where_to_split(args.mini_batch_size, self.ndevices, _add_leftover=True)
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, src_d)):
emb_output_split = self.model.net.Split(
emb_output, lo, split=lp, axis=0
)
"""
# scatter
y = []
for dst_d in range(len(lo)):
src_blob = lo[dst_d]
dst_blob = str(src_blob).replace(
"gpu_" + str(src_d), "gpu_" + str(dst_d), 1
)
if src_blob != dst_blob:
with core.DeviceScope(
core.DeviceOption(workspace.GpuDeviceType, dst_d)
):
blob = self.model.Copy(src_blob, dst_blob)
else:
blob = dst_blob
y.append(blob)
t_list.append(y)
# adjust lists to be ordered per device
x = list(map(lambda x: list(x), zip(*self.bot_l)))
ly = list(map(lambda y: list(y), zip(*t_list)))
# interactions
for d in range(self.ndevices):
on_device = "gpu_" + str(d) + "/"
tag = (on_device + self.tdout, on_device + self.tsout, on_device + self.tint)
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, d)):
self.create_interactions([x[d][-1]], ly[d], self.model, tag)
# replicate mlp (data parallelism)
tag = (self.ttop, self.tint, self.tout)
self.top_l, self.top_w = self.create_mlp(self.ln_top, self.sigmoid_top,
self.model, tag)
# debug prints
# print(self.model.net.Proto(),end='\n')
# sys.exit("ERROR: debugging")
# setup the last output variable
self.last_output = self.top_l[-1]
def __init__(
self,
m_spa,
ln_emb,
ln_bot,
ln_top,
arch_interaction_op,
arch_interaction_itself=False,
sigmoid_bot=-1,
sigmoid_top=-1,
save_onnx=False,
model=None,
test_net=None,
tag=None,
ndevices=-1,
forward_ops=True,
enable_prof=False,
weighted_pooling=None,
emb_optimizer="sgd"
):
super(DLRM_Net, self).__init__()
# init model
if model is None:
global_init_opt = ["caffe2", "--caffe2_log_level=0"]
if enable_prof:
global_init_opt += [
"--logtostderr=0",
"--log_dir=$HOME",
"--caffe2_logging_print_net_summary=1",
]
workspace.GlobalInit(global_init_opt)
self.set_tags()
self.model = model_helper.ModelHelper(name="DLRM", init_params=True)
self.test_net = None
else:
# WARNING: assume that workspace and tags have been initialized elsewhere
self.set_tags(tag[0], tag[1], tag[2], tag[3], tag[4], tag[5], tag[6],
tag[7], tag[8], tag[9])
self.model = model
self.test_net = test_net
# save arguments
self.m_spa = m_spa
self.ln_emb = ln_emb
self.ln_bot = ln_bot
self.ln_top = ln_top
self.arch_interaction_op = arch_interaction_op
self.arch_interaction_itself = arch_interaction_itself
self.sigmoid_bot = sigmoid_bot
self.sigmoid_top = sigmoid_top
self.save_onnx = save_onnx
self.ndevices = ndevices
self.emb_optimizer = emb_optimizer
if weighted_pooling is not None and weighted_pooling != "fixed":
self.weighted_pooling = "learned"
else:
self.weighted_pooling = weighted_pooling
# onnx types and shapes dictionary
if self.save_onnx:
self.onnx_tsd = {}
# create forward operators
if forward_ops:
if self.ndevices <= 1:
return self.create_sequential_forward_ops()
else:
return self.create_parallel_forward_ops()
def set_tags(
self,
_tag_layer_top_mlp="top",
_tag_layer_bot_mlp="bot",
_tag_layer_embedding="emb",
_tag_feature_dense_in="dense_in",
_tag_feature_dense_out="dense_out",
_tag_feature_sparse_in="sparse_in",
_tag_feature_sparse_out="sparse_out",
_tag_interaction="interaction",
_tag_dense_output="prob_click",
_tag_dense_target="target",
):
# layer tags
self.ttop = _tag_layer_top_mlp
self.tbot = _tag_layer_bot_mlp
self.temb = _tag_layer_embedding
# dense feature tags
self.tdin = _tag_feature_dense_in
self.tdout = _tag_feature_dense_out
# sparse feature tags
self.tsin = _tag_feature_sparse_in
self.tsout = _tag_feature_sparse_out
# output and target tags
self.tint = _tag_interaction
self.ttar = _tag_dense_target
self.tout = _tag_dense_output
def parameters(self):
return self.model
def get_loss(self):
return self.FetchBlobWrapper(self.loss, reduce_across="add")
def get_output(self):
return self.FetchBlobWrapper(self.last_output, reduce_across="concat")
def create(self, X, S_lengths, S_indices, T):
self.create_input(X, S_lengths, S_indices, T)
self.create_model(X, S_lengths, S_indices, T)
def create_input(self, X, S_lengths, S_indices, T):
# feed input data to blobs
self.FeedBlobWrapper(self.tdin, X, split=True)
# save the blob shapes for latter (only needed if onnx is requested)
if self.save_onnx:
self.onnx_tsd[self.tdin] = (onnx.TensorProto.FLOAT, X.shape)
for i in range(len(self.emb_l)):
# select device
if self.ndevices > 1:
d = i % self.ndevices
else:
d = -1
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
len_s = on_device + self.temb + ":::" + "sls" + str(i) + "_l"
ind_s = on_device + self.temb + ":::" + "sls" + str(i) + "_i"
self.FeedBlobWrapper(len_s, np.array(S_lengths[i]), False, device_id=d)
self.FeedBlobWrapper(ind_s, np.array(S_indices[i]), False, device_id=d)
# save the blob shapes for latter (only needed if onnx is requested)
if self.save_onnx:
lshape = (len(S_lengths[i]),) # =args.mini_batch_size
ishape = (len(S_indices[i]),)
self.onnx_tsd[len_s] = (onnx.TensorProto.INT32, lshape)
self.onnx_tsd[ind_s] = (onnx.TensorProto.INT32, ishape)
# feed target data to blobs
if T is not None:
zeros_fp32 = np.zeros(T.shape).astype(np.float32)
self.FeedBlobWrapper(self.ttar, zeros_fp32, split=True)
# save the blob shapes for latter (only needed if onnx is requested)
if self.save_onnx:
self.onnx_tsd[self.ttar] = (onnx.TensorProto.FLOAT, T.shape)
def create_model(self, X, S_lengths, S_indices, T):
#setup tril indices for the interactions
offset = 1 if self.arch_interaction_itself else 0
num_fea = len(self.emb_l) + 1
tril_indices = np.array([j + i * num_fea
for i in range(num_fea) for j in range(i + offset)])
self.FeedBlobWrapper(self.tint + "_tril_indices", tril_indices)
# create compute graph
if T is not None:
# WARNING: RunNetOnce call is needed only if we use brew and ConstantFill.
# We could use direct calls to self.model functions above to avoid it
workspace.RunNetOnce(self.model.param_init_net)
workspace.CreateNet(self.model.net)
if self.test_net is not None:
workspace.CreateNet(self.test_net)
def run(self, X, S_lengths, S_indices, T, test_net=False, enable_prof=False):
# feed input data to blobs
# dense features
self.FeedBlobWrapper(self.tdin, X, split=True)
# sparse features
for i in range(len(self.emb_l)):
# select device
if self.ndevices > 1:
d = i % self.ndevices
else:
d = -1
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
len_s = on_device + self.temb + ":::" + "sls" + str(i) + "_l"
ind_s = on_device + self.temb + ":::" + "sls" + str(i) + "_i"
self.FeedBlobWrapper(len_s, np.array(S_lengths[i]), False, device_id=d)
self.FeedBlobWrapper(ind_s, np.array(S_indices[i]), False, device_id=d)
# feed target data to blobs if needed
if T is not None:
self.FeedBlobWrapper(self.ttar, T, split=True)
# execute compute graph
if test_net:
workspace.RunNet(self.test_net)
else:
if enable_prof:
workspace.C.benchmark_net(self.model.net.Name(), 0, 1, True)
else:
workspace.RunNet(self.model.net)
# debug prints
# print("intermediate")
# print(self.FetchBlobWrapper(self.bot_l[-1]))
# for tag_emb in self.emb_l:
# print(self.FetchBlobWrapper(tag_emb))
# print(self.FetchBlobWrapper(self.tint))
def MSEloss(self, scale=1.0):
# add MSEloss to the model
self.AddLayerWrapper(self.model.SquaredL2Distance, [self.tout, self.ttar], "sd")
self.AddLayerWrapper(self.model.Scale, "sd", "sd2", scale=2.0 * scale)
# WARNING: "loss" is a special tag and should not be changed
self.loss = self.AddLayerWrapper(self.model.AveragedLoss, "sd2", "loss")
def BCEloss(self, scale=1.0, threshold=0.0):
# add BCEloss to the mode
if 0.0 < threshold and threshold < 1.0:
self.AddLayerWrapper(self.model.Clip, self.tout, "tout_c",
min=threshold, max=(1.0 - threshold))
self.AddLayerWrapper(self.model.MakeTwoClass, "tout_c", "tout_2c")
else:
self.AddLayerWrapper(self.model.MakeTwoClass, self.tout, "tout_2c")
self.AddLayerWrapper(self.model.LabelCrossEntropy, ["tout_2c", self.ttar], "sd")
# WARNING: "loss" is a special tag and should not be changed
if scale == 1.0:
self.loss = self.AddLayerWrapper(self.model.AveragedLoss, "sd", "loss")
else:
self.AddLayerWrapper(self.model.Scale, "sd", "sd2", scale=scale)
self.loss = self.AddLayerWrapper(self.model.AveragedLoss, "sd2", "loss")
def sgd_optimizer(self, learning_rate,
T=None, _gradientMap=None, sync_dense_params=True):
# create one, it and lr tags (or use them if already present)
if T is not None:
(tag_one, tag_it, tag_lr) = T
else:
(tag_one, tag_it, tag_lr) = ("const_one", "optim_it", "optim_lr")
# approach 1: feed values directly
# self.FeedBlobWrapper(tag_one, np.ones(1).astype(np.float32))
# self.FeedBlobWrapper(tag_it, np.zeros(1).astype(np.int64))
# it = self.AddLayerWrapper(self.model.Iter, tag_it, tag_it)
# lr = self.AddLayerWrapper(self.model.LearningRate, tag_it, tag_lr,
# base_lr=-1 * learning_rate, policy="fixed")
# approach 2: use brew
self.AddLayerWrapper(self.model.param_init_net.ConstantFill,
[], tag_one, shape=[1], value=1.0)
self.AddLayerWrapper(brew.iter, self.model, tag_it)
self.AddLayerWrapper(self.model.LearningRate, tag_it, tag_lr,
base_lr=-1 * learning_rate, policy="fixed")
# save the blob shapes for latter (only needed if onnx is requested)
if self.save_onnx:
self.onnx_tsd[tag_one] = (onnx.TensorProto.FLOAT, (1,))
self.onnx_tsd[tag_it] = (onnx.TensorProto.INT64, (1,))
# create gradient maps (or use them if already present)
if _gradientMap is not None:
self.gradientMap = _gradientMap
else:
if self.loss.__class__ == list:
self.gradientMap = self.model.AddGradientOperators(self.loss)
else:
self.gradientMap = self.model.AddGradientOperators([self.loss])
# update weights
# approach 1: builtin function
# optimizer.build_sgd(self.model, base_learning_rate=learning_rate)
# approach 2: custom code
# top MLP weight and bias
for w in self.top_w:
# allreduce across devices if needed
if sync_dense_params and self.ndevices > 1:
grad_blobs = [
self.gradientMap["gpu_{}/".format(d) + w]
for d in range(self.ndevices)
]
self.model.NCCLAllreduce(grad_blobs, grad_blobs)
# update weights
self.AddLayerWrapper(self.model.WeightedSum,
[w, tag_one, "", tag_lr], w, reset_grad=True)
# bottom MLP weight and bias
for w in self.bot_w:
# allreduce across devices if needed
if sync_dense_params and self.ndevices > 1:
grad_blobs = [
self.gradientMap["gpu_{}/".format(d) + w]
for d in range(self.ndevices)
]
self.model.NCCLAllreduce(grad_blobs, grad_blobs)
# update weights
self.AddLayerWrapper(self.model.WeightedSum,
[w, tag_one, "", tag_lr], w, reset_grad=True)
# update embeddings
for i, w in enumerate(self.emb_w):
# select device
if self.ndevices > 1:
d = i % self.ndevices
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
_tag_one = on_device + tag_one
_tag_lr = on_device + tag_lr
# pickup gradient
w_grad = self.gradientMap[w]
# update weights
if self.ndevices > 1:
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, d)):
self.model.ScatterWeightedSum([w, _tag_one, w_grad.indices,
w_grad.values, _tag_lr], w)
else:
self.model.ScatterWeightedSum([w, _tag_one, w_grad.indices,
w_grad.values, _tag_lr], w)
# update per sample weights
if self.weighted_pooling == "learned":
for i, w in enumerate(self.emb_vw):
# select device
if self.ndevices > 1:
d = i % self.ndevices
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
_tag_one = on_device + tag_one
_tag_lr = on_device + tag_lr
# pickup gradient
w_grad = self.gradientMap[w]
# update weights
if self.ndevices > 1:
with core.DeviceScope(
core.DeviceOption(workspace.GpuDeviceType, d)
):
self.model.ScatterWeightedSum(
[w, _tag_one, w_grad.indices,
w_grad.values, _tag_lr], w
)
else:
self.model.ScatterWeightedSum(
[w, _tag_one, w_grad.indices, w_grad.values, _tag_lr], w
)
def adagrad_optimizer(self, learning_rate,
T=None, _gradientMap=None, sync_dense_params=True,
epsilon=1e-10, decay_=0.0, weight_decay_=0.0):
# create one, it and lr tags (or use them if already present)
if T is not None:
(tag_one, tag_it, tag_lr) = T
else:
(tag_one, tag_it, tag_lr) = ("const_one", "optim_it", "optim_lr")
# approach 1: feed values directly
# self.FeedBlobWrapper(tag_one, np.ones(1).astype(np.float32))
# self.FeedBlobWrapper(tag_it, np.zeros(1).astype(np.int64))
# it = self.AddLayerWrapper(self.model.Iter, tag_it, tag_it)
# lr = self.AddLayerWrapper(self.model.LearningRate, tag_it, tag_lr,
# base_lr=-1 * learning_rate, policy="fixed")
# approach 2: use brew
self.AddLayerWrapper(self.model.param_init_net.ConstantFill,
[], tag_one, shape=[1], value=1.0)
self.AddLayerWrapper(brew.iter, self.model, tag_it)
self.AddLayerWrapper(self.model.LearningRate, tag_it, tag_lr,
base_lr=-1 * learning_rate, policy="fixed")
# save the blob shapes for latter (only needed if onnx is requested)
if self.save_onnx:
self.onnx_tsd[tag_one] = (onnx.TensorProto.FLOAT, (1,))
self.onnx_tsd[tag_it] = (onnx.TensorProto.INT64, (1,))
# create gradient maps (or use them if already present)
if _gradientMap is not None:
self.gradientMap = _gradientMap
else:
if self.loss.__class__ == list:
self.gradientMap = self.model.AddGradientOperators(self.loss)
else:
self.gradientMap = self.model.AddGradientOperators([self.loss])
# update weights
# approach 1: builtin function
# optimizer.build_sgd(self.model, base_learning_rate=learning_rate)
# approach 2: custom code
# top MLP weight and bias
for i, w in enumerate(self.top_w):
# allreduce across devices if needed
if sync_dense_params and self.ndevices > 1:
grad_blobs = [
self.gradientMap["gpu_{}/".format(d) + w]
for d in range(self.ndevices)
]
self.model.NCCLAllreduce(grad_blobs, grad_blobs)
# update weights
self.model.Adagrad(
[
w,
"momentum_mlp_top_{}".format(i + 1),
self.gradientMap[w],
tag_lr
],
[w, "momentum_mlp_top_{}".format(i + 1)],
epsilon=epsilon,
decay_=decay_,
weight_decay_=weight_decay_
)
# bottom MLP weight and bias
for i, w in enumerate(self.bot_w):
# allreduce across devices if needed
if sync_dense_params and self.ndevices > 1:
grad_blobs = [
self.gradientMap["gpu_{}/".format(d) + w]
for d in range(self.ndevices)
]
self.model.NCCLAllreduce(grad_blobs, grad_blobs)
# update weights
self.model.Adagrad(
[
w,
"momentum_mlp_bot_{}".format(i + 1),
self.gradientMap[w],
tag_lr
],
[w, "momentum_mlp_bot_{}".format(i + 1)],
epsilon=epsilon,
decay_=decay_,
weight_decay_=weight_decay_
)
# update embeddings
for i, w in enumerate(self.emb_w):
# select device
if self.ndevices > 1:
d = i % self.ndevices
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
_tag_one = on_device + tag_one
_tag_lr = on_device + tag_lr
# pickup gradient
w_grad = self.gradientMap[w]
# update weights
def add_optimizer():
self.model.Unique(
w_grad.indices,
["unique_w_grad_indices", "remapping_w_grad_indices"]
)
self.model.UnsortedSegmentSum(
[w_grad.values, "remapping_w_grad_indices"],
"unique_w_grad_values"
)
if self.emb_optimizer == "adagrad":
self.model.SparseAdagrad(
[
w,
"momentum_emb_{}".format(i),
"unique_w_grad_indices",
"unique_w_grad_values",
_tag_lr
],
[w, "momentum_emb_{}".format(i)],
epsilon=epsilon,
decay_=decay_,
weight_decay_=weight_decay_
)
elif self.emb_optimizer == "rwsadagrad":
self.model.RowWiseSparseAdagrad(
[
w,
"momentum_emb_{}".format(i),
"unique_w_grad_indices",
"unique_w_grad_values",
_tag_lr
],
[w, "momentum_emb_{}".format(i)],
epsilon=epsilon,
decay_=decay_,
weight_decay_=weight_decay_
)
if self.ndevices > 1:
with core.DeviceScope(core.DeviceOption(workspace.GpuDeviceType, d)):
add_optimizer()
else:
add_optimizer()
# update per sample weights
if self.weighted_pooling == "learned":
for i, w in enumerate(self.emb_vw):
# select device
if self.ndevices > 1:
d = i % self.ndevices
# create tags
on_device = "" if self.ndevices <= 1 else "gpu_" + str(d) + "/"
_tag_one = on_device + tag_one
_tag_lr = on_device + tag_lr
# pickup gradient
w_grad = self.gradientMap[w]
# update weights
if self.ndevices > 1:
with core.DeviceScope(
core.DeviceOption(workspace.GpuDeviceType, d)
):
self.model.ScatterWeightedSum(
[w, _tag_one, w_grad.indices,
w_grad.values, _tag_lr], w
)
else:
self.model.ScatterWeightedSum(
[w, _tag_one, w_grad.indices, w_grad.values, _tag_lr], w
)
def print_all(self):
# approach 1: all
print(workspace.Blobs(), end='\n')
for _, l in enumerate(workspace.Blobs()):
print(l)
print(self.FetchBlobWrapper(l))
# approach 2: only summary
# for param in self.model.params:
# self.model.Summarize(param, [], to_file=1)
# self.model.Summarize(self.model.param_to_grad[param], [], to_file=1)
def print_weights(self):
for _, l in enumerate(self.emb_w):
# print(l)
print(self.FetchBlobWrapper(l, False))
if self.weighted_pooling == "learned":
for _, l in enumerate(self.emb_vw):
# print(l)
print(self.FetchBlobWrapper(l, False))
for _, l in enumerate(self.bot_w):
# print(l)
if self.ndevices > 1:
print(self.FetchBlobWrapper(l, False, device_id=0))
else:
print(self.FetchBlobWrapper(l))
for _, l in enumerate(self.top_w):
# print(l)
if self.ndevices > 1:
print(self.FetchBlobWrapper(l, False, device_id=0))
else:
print(self.FetchBlobWrapper(l))
def print_activations(self):
for _, l in enumerate(self.emb_l):
print(l)
print(self.FetchBlobWrapper(l, False))
for _, l in enumerate(self.bot_l):
print(l)
print(self.FetchBlobWrapper(l))
print(self.tint)
print(self.FetchBlobWrapper(self.tint))
for _, l in enumerate(self.top_l):
print(l)
print(self.FetchBlobWrapper(l))
def define_metrics():
metrics = {
'loss': lambda y_true, y_score:
sklearn.metrics.log_loss(
y_true=y_true,
y_pred=y_score,
labels=[0,1]),
'recall': lambda y_true, y_score:
sklearn.metrics.recall_score(
y_true=y_true,
y_pred=np.round(y_score)
),
'precision': lambda y_true, y_score:
sklearn.metrics.precision_score(
y_true=y_true,
y_pred=np.round(y_score)
),
'f1': lambda y_true, y_score:
sklearn.metrics.f1_score(
y_true=y_true,
y_pred=np.round(y_score)
),
'ap': sklearn.metrics.average_precision_score,
'roc_auc': sklearn.metrics.roc_auc_score,
'accuracy': lambda y_true, y_score:
sklearn.metrics.accuracy_score(
y_true=y_true,
y_pred=np.round(y_score)
),
# 'pre_curve' : sklearn.metrics.precision_recall_curve,
# 'roc_curve' : sklearn.metrics.roc_curve,
}
return metrics
def calculate_metrics(targets, scores):
scores = np.concatenate(scores, axis=0)
targets = np.concatenate(targets, axis=0)
metrics = define_metrics()
# print("Compute time for validation metric : ", end="")
# first_it = True
validation_results = {}
for metric_name, metric_function in metrics.items():
# if first_it:
# first_it = False
# else:
# print(", ", end="")
# metric_compute_start = time_wrap(False)
try:
validation_results[metric_name] = metric_function(
targets,
scores
)
except Exception as error :
validation_results[metric_name] = -1
print("{} in calculating {}".format(error, metric_name))
# metric_compute_end = time_wrap(False)
# met_time = metric_compute_end - metric_compute_start
# print("{} {:.4f}".format(metric_name, 1000 * (met_time)),
# end="")
# print(" ms")
return validation_results
if __name__ == "__main__":
### import packages ###
import sys
import argparse
### parse arguments ###
parser = argparse.ArgumentParser(
description="Train Deep Learning Recommendation Model (DLRM)"
)
# model related parameters
parser.add_argument("--arch-sparse-feature-size", type=int, default=2)
parser.add_argument("--arch-embedding-size", type=str, default="4-3-2")
parser.add_argument("--arch-mlp-bot", type=str, default="4-3-2")
parser.add_argument("--arch-mlp-top", type=str, default="4-2-1")
parser.add_argument("--arch-interaction-op", type=str, default="dot")
parser.add_argument("--arch-interaction-itself", action="store_true", default=False)
# activations and loss
parser.add_argument("--activation-function", type=str, default="relu")
parser.add_argument("--loss-function", type=str, default="mse") # or bce
parser.add_argument("--loss-threshold", type=float, default=0.0) # 1.0e-7
parser.add_argument("--round-targets", type=bool, default=False)
parser.add_argument("--weighted-pooling", type=str, default=None)
# data
parser.add_argument("--data-size", type=int, default=1)
parser.add_argument("--num-batches", type=int, default=0)
parser.add_argument("--data-generation", type=str, default="random") # or synthetic or dataset
parser.add_argument("--rand-data-dist", type=str, default="uniform") # uniform or gaussian
parser.add_argument("--rand-data-min", type=float, default=0)
parser.add_argument("--rand-data-max", type=float, default=1)
parser.add_argument("--rand-data-mu", type=float, default=-1)
parser.add_argument("--rand-data-sigma", type=float, default=1)
parser.add_argument("--data-trace-file", type=str, default="./input/dist_emb_j.log")
parser.add_argument("--data-set", type=str, default="kaggle") # or terabyte
parser.add_argument("--raw-data-file", type=str, default="")
parser.add_argument("--processed-data-file", type=str, default="")
parser.add_argument("--data-randomize", type=str, default="total") # or day or none
parser.add_argument("--data-trace-enable-padding", type=bool, default=False)
parser.add_argument("--max-ind-range", type=int, default=-1)
parser.add_argument("--data-sub-sample-rate", type=float, default=0.0) # in [0, 1]
parser.add_argument("--num-indices-per-lookup", type=int, default=10)
parser.add_argument("--num-indices-per-lookup-fixed", type=bool, default=False)
parser.add_argument("--num-workers", type=int, default=0)
parser.add_argument("--memory-map", action="store_true", default=False)
# training
parser.add_argument("--mini-batch-size", type=int, default=1)
parser.add_argument("--nepochs", type=int, default=1)
parser.add_argument("--learning-rate", type=float, default=0.01)
parser.add_argument("--print-precision", type=int, default=5)
parser.add_argument("--numpy-rand-seed", type=int, default=123)
parser.add_argument("--sync-dense-params", type=bool, default=True)
parser.add_argument("--caffe2-net-type", type=str, default="")
parser.add_argument("--optimizer", type=str, default="sgd",
help="""This is the optimizer for embedding tables.""")
parser.add_argument(
"--dataset-multiprocessing",
action="store_true",
default=False,
help="The Kaggle dataset can be multiprocessed in an environment \
with more than 7 CPU cores and more than 20 GB of memory. \n \
The Terabyte dataset can be multiprocessed in an environment \
with more than 24 CPU cores and at least 1 TB of memory.",
)
# inference
parser.add_argument("--inference-only", action="store_true", default=False)
# onnx (or protobuf with shapes)
parser.add_argument("--save-onnx", action="store_true", default=False)
parser.add_argument("--save-proto-types-shapes", action="store_true", default=False)
# gpu
parser.add_argument("--use-gpu", action="store_true", default=False)
# debugging and profiling
parser.add_argument("--print-freq", type=int, default=1)
parser.add_argument("--test-freq", type=int, default=-1)
parser.add_argument("--test-mini-batch-size", type=int, default=-1)
parser.add_argument("--test-num-workers", type=int, default=-1)
parser.add_argument("--print-time", action="store_true", default=False)
parser.add_argument("--debug-mode", action="store_true", default=False)
parser.add_argument("--enable-profiling", action="store_true", default=False)
parser.add_argument("--plot-compute-graph", action="store_true", default=False)
# mlperf logging (disables other output and stops early)
parser.add_argument("--mlperf-logging", action="store_true", default=False)
# stop at target accuracy Kaggle 0.789, Terabyte (sub-sampled=0.875) 0.8107
parser.add_argument("--mlperf-acc-threshold", type=float, default=0.0)
# stop at target AUC Terabyte (no subsampling) 0.8025
parser.add_argument("--mlperf-auc-threshold", type=float, default=0.0)
args = parser.parse_args()
if args.dataset_multiprocessing:
assert float(sys.version[:3]) > 3.7, "The dataset_multiprocessing " + \
"flag is susceptible to a bug in Python 3.7 and under. " + \
"https://github.com/facebookresearch/dlrm/issues/172"
### some basic setup ###
# WARNING: to obtain exactly the same initialization for
# the weights we need to start from the same random seed.
np.random.seed(args.numpy_rand_seed)
np.set_printoptions(precision=args.print_precision)
if (args.test_mini_batch_size < 0):
# if the parameter is not set, use the training batch size
args.test_mini_batch_size = args.mini_batch_size
if (args.test_num_workers < 0):
# if the parameter is not set, use the same parameter for training
args.test_num_workers = args.num_workers
use_gpu = args.use_gpu
if use_gpu:
device_opt = core.DeviceOption(workspace.GpuDeviceType, 0)
ngpus = workspace.NumGpuDevices() # 1
print("Using {} GPU(s)...".format(ngpus))
else:
device_opt = core.DeviceOption(caffe2_pb2.CPU)
print("Using CPU...")
### prepare training data ###
ln_bot = np.fromstring(args.arch_mlp_bot, dtype=int, sep="-")
if args.data_generation == "dataset":
if args.num_workers > 0 or args.test_num_workers > 0:
print("WARNING: non default --num-workers or --test-num-workers options"
+ " are not supported and will be ignored")
if args.mini_batch_size != args.test_mini_batch_size:
print("WARNING: non default ----test-mini-batch-size option"
+ " is not supported and will be ignored")
# input and target from dataset
train_data, train_ld, test_data, test_ld = \
dp.make_criteo_data_and_loaders(
args,
offset_to_length_converter=True,
)
nbatches = args.num_batches if args.num_batches > 0 \
else len(train_ld)
nbatches_test = len(test_ld)
ln_emb = train_data.counts
m_den = train_data.m_den
# enforce maximum limit on number of vectors per embedding
if args.max_ind_range > 0:
ln_emb = np.array(list(map(
lambda x: x if x < args.max_ind_range else args.max_ind_range,
ln_emb
)))
ln_bot[0] = m_den
else:
if args.num_workers > 0 or args.test_num_workers > 0:
print("WARNING: non default --num-workers or --test-num-workers options"
+ " are not supported and will be ignored")
if args.mini_batch_size != args.test_mini_batch_size:
print("WARNING: non default ----test-mini-batch-size option"
+ " is not supported and will be ignored")
# input and target at random
ln_emb = np.fromstring(args.arch_embedding_size, dtype=int, sep="-")
m_den = ln_bot[0]
train_data, train_ld, test_data, test_ld = dp.make_random_data_and_loader(args, ln_emb, m_den, \
offset_to_length_converter=True,
)
nbatches = args.num_batches if args.num_batches > 0 else len(train_ld)
nbatches_test = len(test_ld)
# table_feature_map = {idx : idx for idx in range(len(ln_emb))}
### parse command line arguments ###
m_spa = args.arch_sparse_feature_size
ln_emb = np.asarray(ln_emb)
num_fea = ln_emb.size + 1 # num sparse + num dense features
m_den_out = ln_bot[ln_bot.size - 1]
if args.arch_interaction_op == "dot":
# approach 1: all
# num_int = num_fea * num_fea + m_den_out
# approach 2: unique
if args.arch_interaction_itself:
num_int = (num_fea * (num_fea + 1)) // 2 + m_den_out
else:
num_int = (num_fea * (num_fea - 1)) // 2 + m_den_out
elif args.arch_interaction_op == "cat":
num_int = num_fea * m_den_out
else:
sys.exit("ERROR: --arch-interaction-op="
+ args.arch_interaction_op + " is not supported")
arch_mlp_top_adjusted = str(num_int) + "-" + args.arch_mlp_top
ln_top = np.fromstring(arch_mlp_top_adjusted, dtype=int, sep="-")
# sanity check: feature sizes and mlp dimensions must match
if m_den != ln_bot[0]:
sys.exit("ERROR: arch-dense-feature-size "
+ str(m_den) + " does not match first dim of bottom mlp " + str(ln_bot[0]))
if m_spa != m_den_out:
sys.exit("ERROR: arch-sparse-feature-size "
+ str(m_spa) + " does not match last dim of bottom mlp " + str(m_den_out))
if num_int != ln_top[0]:
sys.exit("ERROR: # of feature interactions "
+ str(num_int) + " does not match first dim of top mlp " + str(ln_top[0]))
# test prints (model arch)
if args.debug_mode:
print("model arch:")
print("mlp top arch " + str(ln_top.size - 1)
+ " layers, with input to output dimensions:")
print(ln_top)
print("# of interactions")
print(num_int)
print("mlp bot arch " + str(ln_bot.size - 1)
+ " layers, with input to output dimensions:")
print(ln_bot)
print("# of features (sparse and dense)")
print(num_fea)
print("dense feature size")
print(m_den)
print("sparse feature size")
print(m_spa)
print("# of embeddings (= # of sparse features) " + str(ln_emb.size)
+ ", with dimensions " + str(m_spa) + "x:")
print(ln_emb)
print("data (inputs and targets):")
for j, inputBatch in enumerate(train_ld):
lX_j, lS_l_j, lS_i_j, lT_j = inputBatch
print("mini-batch: %d" % j)
print(lX_j)
print(lS_l_j)
print(lS_i_j)
print(lT_j)
### construct the neural network specified above ###
# WARNING: to obtain exactly the same initialization for
# the weights we need to start from the same random seed.
# np.random.seed(args.numpy_rand_seed)
ndevices = min(ngpus, args.mini_batch_size, num_fea - 1) if use_gpu else -1
flag_types_shapes = args.save_onnx or args.save_proto_types_shapes
flag_forward_ops = not (use_gpu and ndevices > 1)
with core.DeviceScope(device_opt):
dlrm = DLRM_Net(
m_spa,
ln_emb,
ln_bot,
ln_top,
args.arch_interaction_op,
arch_interaction_itself=args.arch_interaction_itself,
sigmoid_bot=-1,
sigmoid_top=ln_top.size - 1,
save_onnx=flag_types_shapes,
ndevices=ndevices,
# forward_ops = flag_forward_ops
enable_prof=args.enable_profiling,
weighted_pooling=args.weighted_pooling,
emb_optimizer=args.optimizer
)
# load nccl if using multiple devices
if args.sync_dense_params and ndevices > 1:
dyndep.InitOpsLibrary("//caffe2/caffe2/contrib/nccl:nccl_ops")
# set the net type for better performance (dag, async_scheduling, etc)
if args.caffe2_net_type:
dlrm.parameters().net.Proto().type = args.caffe2_net_type
# plot compute graph
if args.plot_compute_graph:
graph = net_drawer.GetPydotGraph(
dlrm.parameters().net,
"dlrm_s_caffe2_graph",
"BT"
)
graph.write_pdf(graph.get_name() + ".pdf")
# test prints
if args.debug_mode:
print("initial parameters (weights and bias):")
dlrm.print_weights()
# add training loss if needed
if not args.inference_only:
with core.DeviceScope(device_opt):
# specify the loss function
nd = 1.0 if dlrm.ndevices <= 1 else 1.0 / dlrm.ndevices # 1
if args.loss_function == "mse":
dlrm.MSEloss(scale=nd)
elif args.loss_function == "bce":
dlrm.BCEloss(scale=nd, threshold=args.loss_threshold)
else:
sys.exit("ERROR: --loss-function=" + args.loss_function
+ " is not supported")
# define test net (as train net without gradients)
dlrm.test_net = core.Net(copy.deepcopy(dlrm.model.net.Proto()))
# specify the optimizer algorithm
if args.optimizer == "sgd":
dlrm.sgd_optimizer(
args.learning_rate, sync_dense_params=args.sync_dense_params
)
elif args.optimizer in ["adagrad", "rwsadagrad"]:
dlrm.adagrad_optimizer(
args.learning_rate, sync_dense_params=args.sync_dense_params
)
else:
sys.exit("""ERROR: Select an optimizer for
embedding tables : 'sgd', 'adagrad',
or 'rwsadagrad' """)
# init/create
X, lS_l, lS_i, T = next(iter(train_ld)) # does not affect the enumerate(train_ld) in the main loop
dlrm.create(X, lS_l, lS_i, T.int())
### main loop ###
best_gA_test = 0
best_auc_test = 0
total_time = 0
total_loss = 0
total_accu = 0
total_iter = 0
total_samp = 0
k = 0
print("time/loss/accuracy (if enabled):")
while k < args.nepochs:
j = 0
for j, inputBatch in enumerate(train_ld):
# forward and backward pass, where the latter runs only
# when gradients and loss have been added to the net
time1 = time.time()
lX_j, lS_l_j, lS_i_j, lT_j = inputBatch
lT_j = lT_j.int() if args.loss_function == "bce" else lT_j
dlrm.run(lX_j, lS_l_j, lS_i_j, lT_j)
time2 = time.time()
total_time += time2 - time1
# compte loss and accuracy
Z = dlrm.get_output() # numpy array
T = lT_j.numpy()
'''
# debug prints
print("output and loss")
print(Z)
print(dlrm.get_loss())
'''
mbs = T.shape[0] # = args.mini_batch_size except maybe for last
A = np.sum((np.round(Z, 0) == T).astype(np.uint8))
total_accu += 0 if args.inference_only else A
total_loss += 0 if args.inference_only else dlrm.get_loss() * mbs
total_iter += 1
total_samp += mbs
# print time, loss and accuracy
should_print = ((j + 1) % args.print_freq == 0) or (j + 1 == nbatches)
should_test = (
(args.test_freq > 0)
and (args.data_generation in ["dataset", "random"])
and (((j + 1) % args.test_freq == 0) or (j + 1 == nbatches))
)
if should_print or should_test:
gT = 1000. * total_time / total_iter if args.print_time else -1
total_time = 0
gA = total_accu / total_samp
total_accu = 0
gL = total_loss / total_samp
total_loss = 0
str_run_type = "inference" if args.inference_only else "training"
print(
"Finished {} it {}/{} of epoch {}, {:.2f} ms/it,".format(
str_run_type, j + 1, nbatches, k, gT
)
+ " loss {:.6f}".format(gL)
)
total_iter = 0
total_samp = 0
# debug prints
# print(Z)
# print(T)
# testing
if should_test and not args.inference_only:
# don't measure training iter time in a test iteration
if args.mlperf_logging:
previous_iteration_time = None
test_accu = 0
test_loss = 0
test_samp = 0
if args.mlperf_logging:
scores = []
targets = []
for i, testBatch in enumerate(test_ld):
# early exit if nbatches was set by the user and was exceeded
if nbatches > 0 and i >= nbatches:
break
# forward pass
lX_test_i, lS_l_test_i, lS_i_test_i, lT_test_i = testBatch
lT_test_i = lT_test_i.int() if args.loss_function == "bce" else lT_test_i
dlrm.run(lX_test_i, lS_l_test_i, lS_i_test_i, lT_test_i, test_net=True)
Z_test = dlrm.get_output()
T_test = lT_test_i.numpy()
if args.mlperf_logging:
scores.append(Z_test)
targets.append(T_test)
else:
# compte loss and accuracy
L_test = dlrm.get_loss()
mbs_test = T_test.shape[0] # = mini_batch_size except last
A_test = np.sum((np.round(Z_test, 0) == T_test).astype(np.uint8))
test_accu += A_test
test_loss += L_test * mbs_test
test_samp += mbs_test
# compute metrics (after test loop has finished)
if args.mlperf_logging:
validation_results = calculate_metrics(targets, scores)
gA_test = validation_results['accuracy']
gL_test = validation_results['loss']
else:
gA_test = test_accu / test_samp
gL_test = test_loss / test_samp
# print metrics
is_best = gA_test > best_gA_test
if is_best:
best_gA_test = gA_test
if args.mlperf_logging:
is_best = validation_results['roc_auc'] > best_auc_test
if is_best:
best_auc_test = validation_results['roc_auc']
print(
"Testing at - {}/{} of epoch {},".format(j + 1, nbatches, k)
+ " loss {:.6f}, recall {:.4f}, precision {:.4f},".format(
validation_results['loss'],
validation_results['recall'],
validation_results['precision']
)
+ " f1 {:.4f}, ap {:.4f},".format(
validation_results['f1'],
validation_results['ap'],
)
+ " auc {:.4f}, best auc {:.4f},".format(
validation_results['roc_auc'],
best_auc_test
)
+ " accuracy {:3.3f} %, best accuracy {:3.3f} %".format(
validation_results['accuracy'] * 100,
best_gA_test * 100
)
)
else:
print(
"Testing at - {}/{} of epoch {},".format(j + 1, nbatches, 0)
+ " loss {:.6f}, accuracy {:3.3f} %, best {:3.3f} %".format(
gL_test, gA_test * 100, best_gA_test * 100
)
)
# check thresholds
if (args.mlperf_logging
and (args.mlperf_acc_threshold > 0)
and (best_gA_test > args.mlperf_acc_threshold)):
print("MLPerf testing accuracy threshold "
+ str(args.mlperf_acc_threshold)
+ " reached, stop training")
break
if (args.mlperf_logging
and (args.mlperf_auc_threshold > 0)
and (best_auc_test > args.mlperf_auc_threshold)):
print("MLPerf testing auc threshold "
+ str(args.mlperf_auc_threshold)
+ " reached, stop training")
break
j += 1 # nbatches
k += 1 # nepochs
# test prints
if not args.inference_only and args.debug_mode:
print("updated parameters (weights and bias):")
dlrm.print_weights()
# build onnx model from caffe2
if args.save_onnx:
pnet = dlrm.parameters().net.Proto()
inet = dlrm.parameters().param_init_net.Proto()
value_info = dlrm.onnx_tsd # None
# debug prints
# print(value_info)
# WARNING: Why Caffe2 to ONNX net transformation currently does not work?
# 1. ONNX does not support SparseLengthsSum operator directly. A workaround
# could be for the Caffe2 ONNX frontend to indirectly map this operator to
# Gather and ReducedSum ONNX operators, following the PyTorch approach.
c2f = caffe2.python.onnx.frontend.Caffe2Frontend()
dlrm_caffe2_onnx = c2f.caffe2_net_to_onnx_model(pnet, inet, value_info)
# check the onnx model
onnx.checker.check_model(dlrm_caffe2_onnx)
# save model to a file
with open("dlrm_s_caffe2.onnx", "w+") as dlrm_caffe2_onnx_file:
dlrm_caffe2_onnx_file.write(str(dlrm_caffe2_onnx))
# build protobuf with types and shapes
if args.save_proto_types_shapes:
# add types and shapes to protobuf
__TYPE_MAPPING = {
onnx.TensorProto.FLOAT: caffe2_pb2.TensorProto.FLOAT,
onnx.TensorProto.UINT8: caffe2_pb2.TensorProto.UINT8,
onnx.TensorProto.INT8: caffe2_pb2.TensorProto.INT8,
onnx.TensorProto.UINT16: caffe2_pb2.TensorProto.UINT16,
onnx.TensorProto.INT16: caffe2_pb2.TensorProto.INT16,
onnx.TensorProto.INT32: caffe2_pb2.TensorProto.INT32,
onnx.TensorProto.INT64: caffe2_pb2.TensorProto.INT64,
onnx.TensorProto.STRING: caffe2_pb2.TensorProto.STRING,
onnx.TensorProto.BOOL: caffe2_pb2.TensorProto.BOOL,
onnx.TensorProto.FLOAT16: caffe2_pb2.TensorProto.FLOAT16,
onnx.TensorProto.DOUBLE: caffe2_pb2.TensorProto.DOUBLE,
}
pnet = dlrm.parameters().net.Proto()
arg = pnet.arg.add()
arg.name = "input_shape_info"
for i in pnet.external_input:
if i in dlrm.onnx_tsd:
onnx_dtype, shape = dlrm.onnx_tsd[i]
t = arg.tensors.add()
t.name = i
t.data_type = __TYPE_MAPPING[onnx_dtype]
t.dims.extend(shape)
else:
print("Warning: we don't have shape/type info for input: {}".format(i))
# debug print
# print(pnet)
# export the protobuf with types and shapes
with open("dlrm_s_caffe2.proto", "w+") as dlrm_s_proto_file:
dlrm_s_proto_file.write(str(pnet))
"""
# export the protobuf with types and shapes as well as weights
# see https://github.com/pytorch/pytorch/issues/9533
#save
net = dlrm.parameters().net
params = dlrm.parameters().params
init_net, predict_net = mobile_exporter.Export(workspace, net, params)
with open("dlrm_s_caffe2.predict", "wb") as dlrm_s_predict_file:
dlrm_s_predict_file.write(predict_net.SerializeToString())
with open("dlrm_s_caffe2.init", "wb") as dlrm_s_init_file:
dlrm_s_init_file.write(init_net.SerializeToString())
#load
net_def = caffe2_pb2.NetDef()
init_def= caffe2_pb2.NetDef()
with open("dlrm_s_caffe2.predict", "rb") as dlrm_s_predict_file:
net_def.ParseFromString(dlrm_s_predict_file.read())
print(net_def)
with open("dlrm_s_caffe2.init", "rb") as dlrm_s_init_file:
init_def.ParseFromString(dlrm_s_init_file.read())
print(init_def)
"""
PyTorch/Recommendation/dlrm/dlrm_s_pytorch.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Description: an implementation of a deep learning recommendation model (DLRM)
# The model input consists of dense and sparse features. The former is a vector
# of floating point values. The latter is a list of sparse indices into
# embedding tables, which consist of vectors of floating point values.
# The selected vectors are passed to mlp networks denoted by triangles,
# in some cases the vectors are interacted through operators (Ops).
#
# output:
# vector of values
# model: |
# /\
# /__\
# |
# _____________________> Op <___________________
# / | \
# /\ /\ /\
# /__\ /__\ ... /__\
# | | |
# | Op Op
# | ____/__\_____ ____/__\____
# | |_Emb_|____|__| ... |_Emb_|__|___|
# input:
# [ dense features ] [sparse indices] , ..., [sparse indices]
#
# More precise definition of model layers:
# 1) fully connected layers of an mlp
# z = f(y)
# y = Wx + b
#
# 2) embedding lookup (for a list of sparse indices p=[p1,...,pk])
# z = Op(e1,...,ek)
# obtain vectors e1=E[:,p1], ..., ek=E[:,pk]
#
# 3) Operator Op can be one of the following
# Sum(e1,...,ek) = e1 + ... + ek
# Dot(e1,...,ek) = [e1'e1, ..., e1'ek, ..., ek'e1, ..., ek'ek]
# Cat(e1,...,ek) = [e1', ..., ek']'
# where ' denotes transpose operation
#
# References:
# [1] Maxim Naumov, Dheevatsa Mudigere, Hao-Jun Michael Shi, Jianyu Huang,
# Narayanan Sundaram, Jongsoo Park, Xiaodong Wang, Udit Gupta, Carole-Jean Wu,
# Alisson G. Azzolini, Dmytro Dzhulgakov, Andrey Mallevich, Ilia Cherniavskii,
# Yinghai Lu, Raghuraman Krishnamoorthi, Ansha Yu, Volodymyr Kondratenko,
# Stephanie Pereira, Xianjie Chen, Wenlin Chen, Vijay Rao, Bill Jia, Liang Xiong,
# Misha Smelyanskiy, "Deep Learning Recommendation Model for Personalization and
# Recommendation Systems", CoRR, arXiv:1906.00091, 2019
from __future__ import absolute_import, division, print_function, unicode_literals
import argparse
# miscellaneous
import builtins
import datetime
import json
import sys
import time
# onnx
# The onnx import causes deprecation warnings every time workers
# are spawned during testing. So, we filter out those warnings.
import warnings
# data generation
import dlrm_data_pytorch as dp
# For distributed run
import extend_distributed as ext_dist
import mlperf_logger
# numpy
import numpy as np
import sklearn.metrics
# pytorch
import torch
import torch.nn as nn
from torch._ops import ops
from torch.autograd.profiler import record_function
from torch.nn.parallel.parallel_apply import parallel_apply
from torch.nn.parallel.replicate import replicate
from torch.nn.parallel.scatter_gather import gather, scatter
from torch.nn.parameter import Parameter
from torch.optim.lr_scheduler import _LRScheduler
import optim.rwsadagrad as RowWiseSparseAdagrad
from torch.utils.tensorboard import SummaryWriter
# mixed-dimension trick
from tricks.md_embedding_bag import PrEmbeddingBag, md_solver
# quotient-remainder trick
from tricks.qr_embedding_bag import QREmbeddingBag
with warnings.catch_warnings():
warnings.filterwarnings("ignore", category=DeprecationWarning)
try:
import onnx
except ImportError as error:
print("Unable to import onnx. ", error)
# from torchviz import make_dot
# import torch.nn.functional as Functional
# from torch.nn.parameter import Parameter
exc = getattr(builtins, "IOError", "FileNotFoundError")
def time_wrap(use_gpu):
if use_gpu:
torch.cuda.synchronize()
return time.time()
def dlrm_wrap(X, lS_o, lS_i, use_gpu, device, ndevices=1):
with record_function("DLRM forward"):
if use_gpu: # .cuda()
# lS_i can be either a list of tensors or a stacked tensor.
# Handle each case below:
if ndevices == 1:
lS_i = (
[S_i.to(device) for S_i in lS_i]
if isinstance(lS_i, list)
else lS_i.to(device)
)
lS_o = (
[S_o.to(device) for S_o in lS_o]
if isinstance(lS_o, list)
else lS_o.to(device)
)
return dlrm(X.to(device), lS_o, lS_i)
def loss_fn_wrap(Z, T, use_gpu, device):
with record_function("DLRM loss compute"):
if args.loss_function == "mse" or args.loss_function == "bce":
return dlrm.loss_fn(Z, T.to(device))
elif args.loss_function == "wbce":
loss_ws_ = dlrm.loss_ws[T.data.view(-1).long()].view_as(T).to(device)
loss_fn_ = dlrm.loss_fn(Z, T.to(device))
loss_sc_ = loss_ws_ * loss_fn_
return loss_sc_.mean()
# The following function is a wrapper to avoid checking this multiple times in th
# loop below.
def unpack_batch(b):
# Experiment with unweighted samples
return b[0], b[1], b[2], b[3], torch.ones(b[3].size()), None
class LRPolicyScheduler(_LRScheduler):
def __init__(self, optimizer, num_warmup_steps, decay_start_step, num_decay_steps):
self.num_warmup_steps = num_warmup_steps
self.decay_start_step = decay_start_step
self.decay_end_step = decay_start_step + num_decay_steps
self.num_decay_steps = num_decay_steps
if self.decay_start_step < self.num_warmup_steps:
sys.exit("Learning rate warmup must finish before the decay starts")
super(LRPolicyScheduler, self).__init__(optimizer)
def get_lr(self):
step_count = self._step_count
if step_count < self.num_warmup_steps:
# warmup
scale = 1.0 - (self.num_warmup_steps - step_count) / self.num_warmup_steps
lr = [base_lr * scale for base_lr in self.base_lrs]
self.last_lr = lr
elif self.decay_start_step <= step_count and step_count < self.decay_end_step:
# decay
decayed_steps = step_count - self.decay_start_step
scale = ((self.num_decay_steps - decayed_steps) / self.num_decay_steps) ** 2
min_lr = 0.0000001
lr = [max(min_lr, base_lr * scale) for base_lr in self.base_lrs]
self.last_lr = lr
else:
if self.num_decay_steps > 0:
# freeze at last, either because we're after decay
# or because we're between warmup and decay
lr = self.last_lr
else:
# do not adjust
lr = self.base_lrs
return lr
### define dlrm in PyTorch ###
class DLRM_Net(nn.Module):
def create_mlp(self, ln, sigmoid_layer):
# build MLP layer by layer
layers = nn.ModuleList()
for i in range(0, ln.size - 1):
n = ln[i]
m = ln[i + 1]
# construct fully connected operator
LL = nn.Linear(int(n), int(m), bias=True)
# initialize the weights
# with torch.no_grad():
# custom Xavier input, output or two-sided fill
mean = 0.0 # std_dev = np.sqrt(variance)
std_dev = np.sqrt(2 / (m + n)) # np.sqrt(1 / m) # np.sqrt(1 / n)
W = np.random.normal(mean, std_dev, size=(m, n)).astype(np.float32)
std_dev = np.sqrt(1 / m) # np.sqrt(2 / (m + 1))
bt = np.random.normal(mean, std_dev, size=m).astype(np.float32)
# approach 1
LL.weight.data = torch.tensor(W, requires_grad=True)
LL.bias.data = torch.tensor(bt, requires_grad=True)
# approach 2
# LL.weight.data.copy_(torch.tensor(W))
# LL.bias.data.copy_(torch.tensor(bt))
# approach 3
# LL.weight = Parameter(torch.tensor(W),requires_grad=True)
# LL.bias = Parameter(torch.tensor(bt),requires_grad=True)
layers.append(LL)
# construct sigmoid or relu operator
if i == sigmoid_layer:
layers.append(nn.Sigmoid())
else:
layers.append(nn.ReLU())
# approach 1: use ModuleList
# return layers
# approach 2: use Sequential container to wrap all layers
return torch.nn.Sequential(*layers)
def create_emb(self, m, ln, weighted_pooling=None):
emb_l = nn.ModuleList()
v_W_l = []
for i in range(0, ln.size):
if ext_dist.my_size > 1:
if i not in self.local_emb_indices:
continue
n = ln[i]
# construct embedding operator
if self.qr_flag and n > self.qr_threshold:
EE = QREmbeddingBag(
n,
m,
self.qr_collisions,
operation=self.qr_operation,
mode="sum",
sparse=True,
)
elif self.md_flag and n > self.md_threshold:
base = max(m)
_m = m[i] if n > self.md_threshold else base
EE = PrEmbeddingBag(n, _m, base)
# use np initialization as below for consistency...
W = np.random.uniform(
low=-np.sqrt(1 / n), high=np.sqrt(1 / n), size=(n, _m)
).astype(np.float32)
EE.embs.weight.data = torch.tensor(W, requires_grad=True)
else:
EE = nn.EmbeddingBag(n, m, mode="sum", sparse=True)
# initialize embeddings
# nn.init.uniform_(EE.weight, a=-np.sqrt(1 / n), b=np.sqrt(1 / n))
W = np.random.uniform(
low=-np.sqrt(1 / n), high=np.sqrt(1 / n), size=(n, m)
).astype(np.float32)
# approach 1
EE.weight.data = torch.tensor(W, requires_grad=True)
# approach 2
# EE.weight.data.copy_(torch.tensor(W))
# approach 3
# EE.weight = Parameter(torch.tensor(W),requires_grad=True)
if weighted_pooling is None:
v_W_l.append(None)
else:
v_W_l.append(torch.ones(n, dtype=torch.float32))
emb_l.append(EE)
return emb_l, v_W_l
def __init__(
self,
m_spa=None,
ln_emb=None,
ln_bot=None,
ln_top=None,
arch_interaction_op=None,
arch_interaction_itself=False,
sigmoid_bot=-1,
sigmoid_top=-1,
sync_dense_params=True,
loss_threshold=0.0,
ndevices=-1,
qr_flag=False,
qr_operation="mult",
qr_collisions=0,
qr_threshold=200,
md_flag=False,
md_threshold=200,
weighted_pooling=None,
loss_function="bce"
):
super(DLRM_Net, self).__init__()
if (
(m_spa is not None)
and (ln_emb is not None)
and (ln_bot is not None)
and (ln_top is not None)
and (arch_interaction_op is not None)
):
# save arguments
self.ndevices = ndevices
self.output_d = 0
self.parallel_model_batch_size = -1
self.parallel_model_is_not_prepared = True
self.arch_interaction_op = arch_interaction_op
self.arch_interaction_itself = arch_interaction_itself
self.sync_dense_params = sync_dense_params
self.loss_threshold = loss_threshold
self.loss_function=loss_function
if weighted_pooling is not None and weighted_pooling != "fixed":
self.weighted_pooling = "learned"
else:
self.weighted_pooling = weighted_pooling
# create variables for QR embedding if applicable
self.qr_flag = qr_flag
if self.qr_flag:
self.qr_collisions = qr_collisions
self.qr_operation = qr_operation
self.qr_threshold = qr_threshold
# create variables for MD embedding if applicable
self.md_flag = md_flag
if self.md_flag:
self.md_threshold = md_threshold
# If running distributed, get local slice of embedding tables
if ext_dist.my_size > 1:
n_emb = len(ln_emb)
if n_emb < ext_dist.my_size:
sys.exit(
"only (%d) sparse features for (%d) devices, table partitions will fail"
% (n_emb, ext_dist.my_size)
)
self.n_global_emb = n_emb
self.n_local_emb, self.n_emb_per_rank = ext_dist.get_split_lengths(
n_emb
)
self.local_emb_slice = ext_dist.get_my_slice(n_emb)
self.local_emb_indices = list(range(n_emb))[self.local_emb_slice]
# create operators
if ndevices <= 1:
self.emb_l, w_list = self.create_emb(m_spa, ln_emb, weighted_pooling)
if self.weighted_pooling == "learned":
self.v_W_l = nn.ParameterList()
for w in w_list:
self.v_W_l.append(Parameter(w))
else:
self.v_W_l = w_list
self.bot_l = self.create_mlp(ln_bot, sigmoid_bot)
self.top_l = self.create_mlp(ln_top, sigmoid_top)
# quantization
self.quantize_emb = False
self.emb_l_q = []
self.quantize_bits = 32
# specify the loss function
if self.loss_function == "mse":
self.loss_fn = torch.nn.MSELoss(reduction="mean")
elif self.loss_function == "bce":
self.loss_fn = torch.nn.BCELoss(reduction="mean")
elif self.loss_function == "wbce":
self.loss_ws = torch.tensor(
np.fromstring(args.loss_weights, dtype=float, sep="-")
)
self.loss_fn = torch.nn.BCELoss(reduction="none")
else:
sys.exit(
"ERROR: --loss-function=" + self.loss_function + " is not supported"
)
def apply_mlp(self, x, layers):
# approach 1: use ModuleList
# for layer in layers:
# x = layer(x)
# return x
# approach 2: use Sequential container to wrap all layers
return layers(x)
def apply_emb(self, lS_o, lS_i, emb_l, v_W_l):
# WARNING: notice that we are processing the batch at once. We implicitly
# assume that the data is laid out such that:
# 1. each embedding is indexed with a group of sparse indices,
# corresponding to a single lookup
# 2. for each embedding the lookups are further organized into a batch
# 3. for a list of embedding tables there is a list of batched lookups
ly = []
for k, sparse_index_group_batch in enumerate(lS_i):
sparse_offset_group_batch = lS_o[k]
# embedding lookup
# We are using EmbeddingBag, which implicitly uses sum operator.
# The embeddings are represented as tall matrices, with sum
# happening vertically across 0 axis, resulting in a row vector
# E = emb_l[k]
if v_W_l[k] is not None:
per_sample_weights = v_W_l[k].gather(0, sparse_index_group_batch)
else:
per_sample_weights = None
if self.quantize_emb:
s1 = self.emb_l_q[k].element_size() * self.emb_l_q[k].nelement()
s2 = self.emb_l_q[k].element_size() * self.emb_l_q[k].nelement()
print("quantized emb sizes:", s1, s2)
if self.quantize_bits == 4:
QV = ops.quantized.embedding_bag_4bit_rowwise_offsets(
self.emb_l_q[k],
sparse_index_group_batch,
sparse_offset_group_batch,
per_sample_weights=per_sample_weights,
)
elif self.quantize_bits == 8:
QV = ops.quantized.embedding_bag_byte_rowwise_offsets(
self.emb_l_q[k],
sparse_index_group_batch,
sparse_offset_group_batch,
per_sample_weights=per_sample_weights,
)
ly.append(QV)
else:
E = emb_l[k]
V = E(
sparse_index_group_batch,
sparse_offset_group_batch,
per_sample_weights=per_sample_weights,
)
ly.append(V)
# print(ly)
return ly
# using quantizing functions from caffe2/aten/src/ATen/native/quantized/cpu
def quantize_embedding(self, bits):
n = len(self.emb_l)
self.emb_l_q = [None] * n
for k in range(n):
if bits == 4:
self.emb_l_q[k] = ops.quantized.embedding_bag_4bit_prepack(
self.emb_l[k].weight
)
elif bits == 8:
self.emb_l_q[k] = ops.quantized.embedding_bag_byte_prepack(
self.emb_l[k].weight
)
else:
return
self.emb_l = None
self.quantize_emb = True
self.quantize_bits = bits
def interact_features(self, x, ly):
if self.arch_interaction_op == "dot":
# concatenate dense and sparse features
(batch_size, d) = x.shape
T = torch.cat([x] + ly, dim=1).view((batch_size, -1, d))
# perform a dot product
Z = torch.bmm(T, torch.transpose(T, 1, 2))
# append dense feature with the interactions (into a row vector)
# approach 1: all
# Zflat = Z.view((batch_size, -1))
# approach 2: unique
_, ni, nj = Z.shape
# approach 1: tril_indices
# offset = 0 if self.arch_interaction_itself else -1
# li, lj = torch.tril_indices(ni, nj, offset=offset)
# approach 2: custom
offset = 1 if self.arch_interaction_itself else 0
li = torch.tensor([i for i in range(ni) for j in range(i + offset)])
lj = torch.tensor([j for i in range(nj) for j in range(i + offset)])
Zflat = Z[:, li, lj]
# concatenate dense features and interactions
R = torch.cat([x] + [Zflat], dim=1)
elif self.arch_interaction_op == "cat":
# concatenation features (into a row vector)
R = torch.cat([x] + ly, dim=1)
else:
sys.exit(
"ERROR: --arch-interaction-op="
+ self.arch_interaction_op
+ " is not supported"
)
return R
def forward(self, dense_x, lS_o, lS_i):
if ext_dist.my_size > 1:
# multi-node multi-device run
return self.distributed_forward(dense_x, lS_o, lS_i)
elif self.ndevices <= 1:
# single device run
return self.sequential_forward(dense_x, lS_o, lS_i)
else:
# single-node multi-device run
return self.parallel_forward(dense_x, lS_o, lS_i)
def distributed_forward(self, dense_x, lS_o, lS_i):
batch_size = dense_x.size()[0]
# WARNING: # of ranks must be <= batch size in distributed_forward call
if batch_size < ext_dist.my_size:
sys.exit(
"ERROR: batch_size (%d) must be larger than number of ranks (%d)"
% (batch_size, ext_dist.my_size)
)
if batch_size % ext_dist.my_size != 0:
sys.exit(
"ERROR: batch_size %d can not split across %d ranks evenly"
% (batch_size, ext_dist.my_size)
)
dense_x = dense_x[ext_dist.get_my_slice(batch_size)]
lS_o = lS_o[self.local_emb_slice]
lS_i = lS_i[self.local_emb_slice]
if (len(self.emb_l) != len(lS_o)) or (len(self.emb_l) != len(lS_i)):
sys.exit(
"ERROR: corrupted model input detected in distributed_forward call"
)
# embeddings
with record_function("DLRM embedding forward"):
ly = self.apply_emb(lS_o, lS_i, self.emb_l, self.v_W_l)
# WARNING: Note that at this point we have the result of the embedding lookup
# for the entire batch on each rank. We would like to obtain partial results
# corresponding to all embedding lookups, but part of the batch on each rank.
# Therefore, matching the distribution of output of bottom mlp, so that both
# could be used for subsequent interactions on each device.
if len(self.emb_l) != len(ly):
sys.exit("ERROR: corrupted intermediate result in distributed_forward call")
a2a_req = ext_dist.alltoall(ly, self.n_emb_per_rank)
with record_function("DLRM bottom nlp forward"):
x = self.apply_mlp(dense_x, self.bot_l)
ly = a2a_req.wait()
ly = list(ly)
# interactions
with record_function("DLRM interaction forward"):
z = self.interact_features(x, ly)
# top mlp
with record_function("DLRM top nlp forward"):
p = self.apply_mlp(z, self.top_l)
# clamp output if needed
if 0.0 < self.loss_threshold and self.loss_threshold < 1.0:
z = torch.clamp(p, min=self.loss_threshold, max=(1.0 - self.loss_threshold))
else:
z = p
return z
def sequential_forward(self, dense_x, lS_o, lS_i):
# process dense features (using bottom mlp), resulting in a row vector
x = self.apply_mlp(dense_x, self.bot_l)
# debug prints
# print("intermediate")
# print(x.detach().cpu().numpy())
# process sparse features(using embeddings), resulting in a list of row vectors
ly = self.apply_emb(lS_o, lS_i, self.emb_l, self.v_W_l)
# for y in ly:
# print(y.detach().cpu().numpy())
# interact features (dense and sparse)
z = self.interact_features(x, ly)
# print(z.detach().cpu().numpy())
# obtain probability of a click (using top mlp)
p = self.apply_mlp(z, self.top_l)
# clamp output if needed
if 0.0 < self.loss_threshold and self.loss_threshold < 1.0:
z = torch.clamp(p, min=self.loss_threshold, max=(1.0 - self.loss_threshold))
else:
z = p
return z
def parallel_forward(self, dense_x, lS_o, lS_i):
### prepare model (overwrite) ###
# WARNING: # of devices must be >= batch size in parallel_forward call
batch_size = dense_x.size()[0]
ndevices = min(self.ndevices, batch_size, len(self.emb_l))
device_ids = range(ndevices)
# WARNING: must redistribute the model if mini-batch size changes(this is common
# for last mini-batch, when # of elements in the dataset/batch size is not even
if self.parallel_model_batch_size != batch_size:
self.parallel_model_is_not_prepared = True
if self.parallel_model_is_not_prepared or self.sync_dense_params:
# replicate mlp (data parallelism)
self.bot_l_replicas = replicate(self.bot_l, device_ids)
self.top_l_replicas = replicate(self.top_l, device_ids)
self.parallel_model_batch_size = batch_size
if self.parallel_model_is_not_prepared:
# distribute embeddings (model parallelism)
t_list = []
w_list = []
for k, emb in enumerate(self.emb_l):
d = torch.device("cuda:" + str(k % ndevices))
t_list.append(emb.to(d))
if self.weighted_pooling == "learned":
w_list.append(Parameter(self.v_W_l[k].to(d)))
elif self.weighted_pooling == "fixed":
w_list.append(self.v_W_l[k].to(d))
else:
w_list.append(None)
self.emb_l = nn.ModuleList(t_list)
if self.weighted_pooling == "learned":
self.v_W_l = nn.ParameterList(w_list)
else:
self.v_W_l = w_list
self.parallel_model_is_not_prepared = False
### prepare input (overwrite) ###
# scatter dense features (data parallelism)
# print(dense_x.device)
dense_x = scatter(dense_x, device_ids, dim=0)
# distribute sparse features (model parallelism)
if (len(self.emb_l) != len(lS_o)) or (len(self.emb_l) != len(lS_i)):
sys.exit("ERROR: corrupted model input detected in parallel_forward call")
t_list = []
i_list = []
for k, _ in enumerate(self.emb_l):
d = torch.device("cuda:" + str(k % ndevices))
t_list.append(lS_o[k].to(d))
i_list.append(lS_i[k].to(d))
lS_o = t_list
lS_i = i_list
### compute results in parallel ###
# bottom mlp
# WARNING: Note that the self.bot_l is a list of bottom mlp modules
# that have been replicated across devices, while dense_x is a tuple of dense
# inputs that has been scattered across devices on the first (batch) dimension.
# The output is a list of tensors scattered across devices according to the
# distribution of dense_x.
x = parallel_apply(self.bot_l_replicas, dense_x, None, device_ids)
# debug prints
# print(x)
# embeddings
ly = self.apply_emb(lS_o, lS_i, self.emb_l, self.v_W_l)
# debug prints
# print(ly)
# butterfly shuffle (implemented inefficiently for now)
# WARNING: Note that at this point we have the result of the embedding lookup
# for the entire batch on each device. We would like to obtain partial results
# corresponding to all embedding lookups, but part of the batch on each device.
# Therefore, matching the distribution of output of bottom mlp, so that both
# could be used for subsequent interactions on each device.
if len(self.emb_l) != len(ly):
sys.exit("ERROR: corrupted intermediate result in parallel_forward call")
t_list = []
for k, _ in enumerate(self.emb_l):
d = torch.device("cuda:" + str(k % ndevices))
y = scatter(ly[k], device_ids, dim=0)
t_list.append(y)
# adjust the list to be ordered per device
ly = list(map(lambda y: list(y), zip(*t_list)))
# debug prints
# print(ly)
# interactions
z = []
for k in range(ndevices):
zk = self.interact_features(x[k], ly[k])
z.append(zk)
# debug prints
# print(z)
# top mlp
# WARNING: Note that the self.top_l is a list of top mlp modules that
# have been replicated across devices, while z is a list of interaction results
# that by construction are scattered across devices on the first (batch) dim.
# The output is a list of tensors scattered across devices according to the
# distribution of z.
p = parallel_apply(self.top_l_replicas, z, None, device_ids)
### gather the distributed results ###
p0 = gather(p, self.output_d, dim=0)
# clamp output if needed
if 0.0 < self.loss_threshold and self.loss_threshold < 1.0:
z0 = torch.clamp(
p0, min=self.loss_threshold, max=(1.0 - self.loss_threshold)
)
else:
z0 = p0
return z0
def dash_separated_ints(value):
vals = value.split("-")
for val in vals:
try:
int(val)
except ValueError:
raise argparse.ArgumentTypeError(
"%s is not a valid dash separated list of ints" % value
)
return value
def dash_separated_floats(value):
vals = value.split("-")
for val in vals:
try:
float(val)
except ValueError:
raise argparse.ArgumentTypeError(
"%s is not a valid dash separated list of floats" % value
)
return value
def inference(
args,
dlrm,
best_acc_test,
best_auc_test,
test_ld,
device,
use_gpu,
log_iter=-1,
):
test_accu = 0
test_samp = 0
if args.mlperf_logging:
scores = []
targets = []
for i, testBatch in enumerate(test_ld):
# early exit if nbatches was set by the user and was exceeded
if nbatches > 0 and i >= nbatches:
break
X_test, lS_o_test, lS_i_test, T_test, W_test, CBPP_test = unpack_batch(
testBatch
)
# Skip the batch if batch size not multiple of total ranks
if ext_dist.my_size > 1 and X_test.size(0) % ext_dist.my_size != 0:
print("Warning: Skiping the batch %d with size %d" % (i, X_test.size(0)))
continue
# forward pass
Z_test = dlrm_wrap(
X_test,
lS_o_test,
lS_i_test,
use_gpu,
device,
ndevices=ndevices,
)
### gather the distributed results on each rank ###
# For some reason it requires explicit sync before all_gather call if
# tensor is on GPU memory
if Z_test.is_cuda:
torch.cuda.synchronize()
(_, batch_split_lengths) = ext_dist.get_split_lengths(X_test.size(0))
if ext_dist.my_size > 1:
Z_test = ext_dist.all_gather(Z_test, batch_split_lengths)
if args.mlperf_logging:
S_test = Z_test.detach().cpu().numpy() # numpy array
T_test = T_test.detach().cpu().numpy() # numpy array
scores.append(S_test)
targets.append(T_test)
else:
with record_function("DLRM accuracy compute"):
# compute loss and accuracy
S_test = Z_test.detach().cpu().numpy() # numpy array
T_test = T_test.detach().cpu().numpy() # numpy array
mbs_test = T_test.shape[0] # = mini_batch_size except last
A_test = np.sum((np.round(S_test, 0) == T_test).astype(np.uint8))
test_accu += A_test
test_samp += mbs_test
if args.mlperf_logging:
with record_function("DLRM mlperf sklearn metrics compute"):
scores = np.concatenate(scores, axis=0)
targets = np.concatenate(targets, axis=0)
metrics = {
"recall": lambda y_true, y_score: sklearn.metrics.recall_score(
y_true=y_true, y_pred=np.round(y_score)
),
"precision": lambda y_true, y_score: sklearn.metrics.precision_score(
y_true=y_true, y_pred=np.round(y_score)
),
"f1": lambda y_true, y_score: sklearn.metrics.f1_score(
y_true=y_true, y_pred=np.round(y_score)
),
"ap": sklearn.metrics.average_precision_score,
"roc_auc": sklearn.metrics.roc_auc_score,
"accuracy": lambda y_true, y_score: sklearn.metrics.accuracy_score(
y_true=y_true, y_pred=np.round(y_score)
),
}
validation_results = {}
for metric_name, metric_function in metrics.items():
validation_results[metric_name] = metric_function(targets, scores)
writer.add_scalar(
"mlperf-metrics-test/" + metric_name,
validation_results[metric_name],
log_iter,
)
acc_test = validation_results["accuracy"]
else:
acc_test = test_accu / test_samp
writer.add_scalar("Test/Acc", acc_test, log_iter)
model_metrics_dict = {
"nepochs": args.nepochs,
"nbatches": nbatches,
"nbatches_test": nbatches_test,
"state_dict": dlrm.state_dict(),
"test_acc": acc_test,
}
if args.mlperf_logging:
is_best = validation_results["roc_auc"] > best_auc_test
if is_best:
best_auc_test = validation_results["roc_auc"]
model_metrics_dict["test_auc"] = best_auc_test
print(
"recall {:.4f}, precision {:.4f},".format(
validation_results["recall"],
validation_results["precision"],
)
+ " f1 {:.4f}, ap {:.4f},".format(
validation_results["f1"], validation_results["ap"]
)
+ " auc {:.4f}, best auc {:.4f},".format(
validation_results["roc_auc"], best_auc_test
)
+ " accuracy {:3.3f} %, best accuracy {:3.3f} %".format(
validation_results["accuracy"] * 100, best_acc_test * 100
),
flush=True,
)
else:
is_best = acc_test > best_acc_test
if is_best:
best_acc_test = acc_test
print(
" accuracy {:3.3f} %, best {:3.3f} %".format(
acc_test * 100, best_acc_test * 100
),
flush=True,
)
return model_metrics_dict, is_best
def run():
### parse arguments ###
parser = argparse.ArgumentParser(
description="Train Deep Learning Recommendation Model (DLRM)"
)
# model related parameters
parser.add_argument("--arch-sparse-feature-size", type=int, default=2)
parser.add_argument(
"--arch-embedding-size", type=dash_separated_ints, default="4-3-2"
)
# j will be replaced with the table number
parser.add_argument("--arch-mlp-bot", type=dash_separated_ints, default="4-3-2")
parser.add_argument("--arch-mlp-top", type=dash_separated_ints, default="4-2-1")
parser.add_argument(
"--arch-interaction-op", type=str, choices=["dot", "cat"], default="dot"
)
parser.add_argument("--arch-interaction-itself", action="store_true", default=False)
parser.add_argument("--weighted-pooling", type=str, default=None)
# embedding table options
parser.add_argument("--md-flag", action="store_true", default=False)
parser.add_argument("--md-threshold", type=int, default=200)
parser.add_argument("--md-temperature", type=float, default=0.3)
parser.add_argument("--md-round-dims", action="store_true", default=False)
parser.add_argument("--qr-flag", action="store_true", default=False)
parser.add_argument("--qr-threshold", type=int, default=200)
parser.add_argument("--qr-operation", type=str, default="mult")
parser.add_argument("--qr-collisions", type=int, default=4)
# activations and loss
parser.add_argument("--activation-function", type=str, default="relu")
parser.add_argument("--loss-function", type=str, default="mse") # or bce or wbce
parser.add_argument(
"--loss-weights", type=dash_separated_floats, default="1.0-1.0"
) # for wbce
parser.add_argument("--loss-threshold", type=float, default=0.0) # 1.0e-7
parser.add_argument("--round-targets", type=bool, default=False)
# data
parser.add_argument("--data-size", type=int, default=1)
parser.add_argument("--num-batches", type=int, default=0)
parser.add_argument(
"--data-generation", type=str, default="random"
) # synthetic or dataset
parser.add_argument(
"--rand-data-dist", type=str, default="uniform"
) # uniform or gaussian
parser.add_argument("--rand-data-min", type=float, default=0)
parser.add_argument("--rand-data-max", type=float, default=1)
parser.add_argument("--rand-data-mu", type=float, default=-1)
parser.add_argument("--rand-data-sigma", type=float, default=1)
parser.add_argument("--data-trace-file", type=str, default="./input/dist_emb_j.log")
parser.add_argument("--data-set", type=str, default="kaggle") # or terabyte
parser.add_argument("--raw-data-file", type=str, default="")
parser.add_argument("--processed-data-file", type=str, default="")
parser.add_argument("--data-randomize", type=str, default="total") # or day or none
parser.add_argument("--data-trace-enable-padding", type=bool, default=False)
parser.add_argument("--max-ind-range", type=int, default=-1)
parser.add_argument("--data-sub-sample-rate", type=float, default=0.0) # in [0, 1]
parser.add_argument("--num-indices-per-lookup", type=int, default=10)
parser.add_argument("--num-indices-per-lookup-fixed", type=bool, default=False)
parser.add_argument("--num-workers", type=int, default=0)
parser.add_argument("--memory-map", action="store_true", default=False)
# training
parser.add_argument("--mini-batch-size", type=int, default=1)
parser.add_argument("--nepochs", type=int, default=1)
parser.add_argument("--learning-rate", type=float, default=0.01)
parser.add_argument("--print-precision", type=int, default=5)
parser.add_argument("--numpy-rand-seed", type=int, default=123)
parser.add_argument("--sync-dense-params", type=bool, default=True)
parser.add_argument("--optimizer", type=str, default="sgd")
parser.add_argument(
"--dataset-multiprocessing",
action="store_true",
default=False,
help="The Kaggle dataset can be multiprocessed in an environment \
with more than 7 CPU cores and more than 20 GB of memory. \n \
The Terabyte dataset can be multiprocessed in an environment \
with more than 24 CPU cores and at least 1 TB of memory.",
)
# inference
parser.add_argument("--inference-only", action="store_true", default=False)
# quantize
parser.add_argument("--quantize-mlp-with-bit", type=int, default=32)
parser.add_argument("--quantize-emb-with-bit", type=int, default=32)
# onnx
parser.add_argument("--save-onnx", action="store_true", default=False)
# gpu
parser.add_argument("--use-gpu", action="store_true", default=False)
# distributed
parser.add_argument("--local_rank", type=int, default=-1)
parser.add_argument("--dist-backend", type=str, default="")
# debugging and profiling
parser.add_argument("--print-freq", type=int, default=1)
parser.add_argument("--test-freq", type=int, default=-1)
parser.add_argument("--test-mini-batch-size", type=int, default=-1)
parser.add_argument("--test-num-workers", type=int, default=-1)
parser.add_argument("--print-time", action="store_true", default=False)
parser.add_argument("--print-wall-time", action="store_true", default=False)
parser.add_argument("--debug-mode", action="store_true", default=False)
parser.add_argument("--enable-profiling", action="store_true", default=False)
parser.add_argument("--plot-compute-graph", action="store_true", default=False)
parser.add_argument("--tensor-board-filename", type=str, default="run_kaggle_pt")
# store/load model
parser.add_argument("--save-model", type=str, default="")
parser.add_argument("--load-model", type=str, default="")
# mlperf logging (disables other output and stops early)
parser.add_argument("--mlperf-logging", action="store_true", default=False)
# stop at target accuracy Kaggle 0.789, Terabyte (sub-sampled=0.875) 0.8107
parser.add_argument("--mlperf-acc-threshold", type=float, default=0.0)
# stop at target AUC Terabyte (no subsampling) 0.8025
parser.add_argument("--mlperf-auc-threshold", type=float, default=0.0)
parser.add_argument("--mlperf-bin-loader", action="store_true", default=False)
parser.add_argument("--mlperf-bin-shuffle", action="store_true", default=False)
# mlperf gradient accumulation iterations
parser.add_argument("--mlperf-grad-accum-iter", type=int, default=1)
# LR policy
parser.add_argument("--lr-num-warmup-steps", type=int, default=0)
parser.add_argument("--lr-decay-start-step", type=int, default=0)
parser.add_argument("--lr-num-decay-steps", type=int, default=0)
global args
global nbatches
global nbatches_test
global writer
args = parser.parse_args()
if args.dataset_multiprocessing:
assert float(sys.version[:3]) > 3.7, "The dataset_multiprocessing " + \
"flag is susceptible to a bug in Python 3.7 and under. " + \
"https://github.com/facebookresearch/dlrm/issues/172"
if args.mlperf_logging:
mlperf_logger.log_event(key=mlperf_logger.constants.CACHE_CLEAR, value=True)
mlperf_logger.log_start(
key=mlperf_logger.constants.INIT_START, log_all_ranks=True
)
if args.weighted_pooling is not None:
if args.qr_flag:
sys.exit("ERROR: quotient remainder with weighted pooling is not supported")
if args.md_flag:
sys.exit("ERROR: mixed dimensions with weighted pooling is not supported")
if args.quantize_emb_with_bit in [4, 8]:
if args.qr_flag:
sys.exit(
"ERROR: 4 and 8-bit quantization with quotient remainder is not supported"
)
if args.md_flag:
sys.exit(
"ERROR: 4 and 8-bit quantization with mixed dimensions is not supported"
)
if args.use_gpu:
sys.exit(
"ERROR: 4 and 8-bit quantization on GPU is not supported"
)
### some basic setup ###
np.random.seed(args.numpy_rand_seed)
np.set_printoptions(precision=args.print_precision)
torch.set_printoptions(precision=args.print_precision)
torch.manual_seed(args.numpy_rand_seed)
if args.test_mini_batch_size < 0:
# if the parameter is not set, use the training batch size
args.test_mini_batch_size = args.mini_batch_size
if args.test_num_workers < 0:
# if the parameter is not set, use the same parameter for training
args.test_num_workers = args.num_workers
use_gpu = args.use_gpu and torch.cuda.is_available()
if not args.debug_mode:
ext_dist.init_distributed(local_rank=args.local_rank, use_gpu=use_gpu, backend=args.dist_backend)
if use_gpu:
torch.cuda.manual_seed_all(args.numpy_rand_seed)
torch.backends.cudnn.deterministic = True
if ext_dist.my_size > 1:
ngpus = 1
device = torch.device("cuda", ext_dist.my_local_rank)
else:
ngpus = torch.cuda.device_count()
device = torch.device("cuda", 0)
print("Using {} GPU(s)...".format(ngpus))
else:
device = torch.device("cpu")
print("Using CPU...")
### prepare training data ###
ln_bot = np.fromstring(args.arch_mlp_bot, dtype=int, sep="-")
# input data
if args.mlperf_logging:
mlperf_logger.barrier()
mlperf_logger.log_end(key=mlperf_logger.constants.INIT_STOP)
mlperf_logger.barrier()
mlperf_logger.log_start(key=mlperf_logger.constants.RUN_START)
mlperf_logger.barrier()
if args.data_generation == "dataset":
train_data, train_ld, test_data, test_ld = dp.make_criteo_data_and_loaders(args)
table_feature_map = {idx: idx for idx in range(len(train_data.counts))}
nbatches = args.num_batches if args.num_batches > 0 else len(train_ld)
nbatches_test = len(test_ld)
ln_emb = train_data.counts
# enforce maximum limit on number of vectors per embedding
if args.max_ind_range > 0:
ln_emb = np.array(
list(
map(
lambda x: x if x < args.max_ind_range else args.max_ind_range,
ln_emb,
)
)
)
else:
ln_emb = np.array(ln_emb)
m_den = train_data.m_den
ln_bot[0] = m_den
else:
# input and target at random
ln_emb = np.fromstring(args.arch_embedding_size, dtype=int, sep="-")
m_den = ln_bot[0]
train_data, train_ld, test_data, test_ld = dp.make_random_data_and_loader(args, ln_emb, m_den)
nbatches = args.num_batches if args.num_batches > 0 else len(train_ld)
nbatches_test = len(test_ld)
args.ln_emb = ln_emb.tolist()
if args.mlperf_logging:
print("command line args: ", json.dumps(vars(args)))
### parse command line arguments ###
m_spa = args.arch_sparse_feature_size
ln_emb = np.asarray(ln_emb)
num_fea = ln_emb.size + 1 # num sparse + num dense features
m_den_out = ln_bot[ln_bot.size - 1]
if args.arch_interaction_op == "dot":
# approach 1: all
# num_int = num_fea * num_fea + m_den_out
# approach 2: unique
if args.arch_interaction_itself:
num_int = (num_fea * (num_fea + 1)) // 2 + m_den_out
else:
num_int = (num_fea * (num_fea - 1)) // 2 + m_den_out
elif args.arch_interaction_op == "cat":
num_int = num_fea * m_den_out
else:
sys.exit(
"ERROR: --arch-interaction-op="
+ args.arch_interaction_op
+ " is not supported"
)
arch_mlp_top_adjusted = str(num_int) + "-" + args.arch_mlp_top
ln_top = np.fromstring(arch_mlp_top_adjusted, dtype=int, sep="-")
# sanity check: feature sizes and mlp dimensions must match
if m_den != ln_bot[0]:
sys.exit(
"ERROR: arch-dense-feature-size "
+ str(m_den)
+ " does not match first dim of bottom mlp "
+ str(ln_bot[0])
)
if args.qr_flag:
if args.qr_operation == "concat" and 2 * m_spa != m_den_out:
sys.exit(
"ERROR: 2 arch-sparse-feature-size "
+ str(2 * m_spa)
+ " does not match last dim of bottom mlp "
+ str(m_den_out)
+ " (note that the last dim of bottom mlp must be 2x the embedding dim)"
)
if args.qr_operation != "concat" and m_spa != m_den_out:
sys.exit(
"ERROR: arch-sparse-feature-size "
+ str(m_spa)
+ " does not match last dim of bottom mlp "
+ str(m_den_out)
)
else:
if m_spa != m_den_out:
sys.exit(
"ERROR: arch-sparse-feature-size "
+ str(m_spa)
+ " does not match last dim of bottom mlp "
+ str(m_den_out)
)
if num_int != ln_top[0]:
sys.exit(
"ERROR: # of feature interactions "
+ str(num_int)
+ " does not match first dimension of top mlp "
+ str(ln_top[0])
)
# assign mixed dimensions if applicable
if args.md_flag:
m_spa = md_solver(
torch.tensor(ln_emb),
args.md_temperature, # alpha
d0=m_spa,
round_dim=args.md_round_dims,
).tolist()
# test prints (model arch)
if args.debug_mode:
print("model arch:")
print(
"mlp top arch "
+ str(ln_top.size - 1)
+ " layers, with input to output dimensions:"
)
print(ln_top)
print("# of interactions")
print(num_int)
print(
"mlp bot arch "
+ str(ln_bot.size - 1)
+ " layers, with input to output dimensions:"
)
print(ln_bot)
print("# of features (sparse and dense)")
print(num_fea)
print("dense feature size")
print(m_den)
print("sparse feature size")
print(m_spa)
print(
"# of embeddings (= # of sparse features) "
+ str(ln_emb.size)
+ ", with dimensions "
+ str(m_spa)
+ "x:"
)
print(ln_emb)
print("data (inputs and targets):")
for j, inputBatch in enumerate(train_ld):
X, lS_o, lS_i, T, W, CBPP = unpack_batch(inputBatch)
torch.set_printoptions(precision=4)
# early exit if nbatches was set by the user and has been exceeded
if nbatches > 0 and j >= nbatches:
break
print("mini-batch: %d" % j)
print(X.detach().cpu())
# transform offsets to lengths when printing
print(
torch.IntTensor(
[
np.diff(
S_o.detach().cpu().tolist() + list(lS_i[i].shape)
).tolist()
for i, S_o in enumerate(lS_o)
]
)
)
print([S_i.detach().cpu() for S_i in lS_i])
print(T.detach().cpu())
global ndevices
ndevices = min(ngpus, args.mini_batch_size, num_fea - 1) if use_gpu else -1
### construct the neural network specified above ###
# WARNING: to obtain exactly the same initialization for
# the weights we need to start from the same random seed.
# np.random.seed(args.numpy_rand_seed)
global dlrm
dlrm = DLRM_Net(
m_spa,
ln_emb,
ln_bot,
ln_top,
arch_interaction_op=args.arch_interaction_op,
arch_interaction_itself=args.arch_interaction_itself,
sigmoid_bot=-1,
sigmoid_top=ln_top.size - 2,
sync_dense_params=args.sync_dense_params,
loss_threshold=args.loss_threshold,
ndevices=ndevices,
qr_flag=args.qr_flag,
qr_operation=args.qr_operation,
qr_collisions=args.qr_collisions,
qr_threshold=args.qr_threshold,
md_flag=args.md_flag,
md_threshold=args.md_threshold,
weighted_pooling=args.weighted_pooling,
loss_function=args.loss_function
)
# test prints
if args.debug_mode:
print("initial parameters (weights and bias):")
for param in dlrm.parameters():
print(param.detach().cpu().numpy())
# print(dlrm)
if use_gpu:
# Custom Model-Data Parallel
# the mlps are replicated and use data parallelism, while
# the embeddings are distributed and use model parallelism
dlrm = dlrm.to(device) # .cuda()
if dlrm.ndevices > 1:
dlrm.emb_l, dlrm.v_W_l = dlrm.create_emb(
m_spa, ln_emb, args.weighted_pooling
)
else:
if dlrm.weighted_pooling == "fixed":
for k, w in enumerate(dlrm.v_W_l):
dlrm.v_W_l[k] = w.cuda()
# distribute data parallel mlps
if ext_dist.my_size > 1:
if use_gpu:
device_ids = [ext_dist.my_local_rank]
dlrm.bot_l = ext_dist.DDP(dlrm.bot_l, device_ids=device_ids)
dlrm.top_l = ext_dist.DDP(dlrm.top_l, device_ids=device_ids)
else:
dlrm.bot_l = ext_dist.DDP(dlrm.bot_l)
dlrm.top_l = ext_dist.DDP(dlrm.top_l)
if not args.inference_only:
if use_gpu and args.optimizer in ["rwsadagrad", "adagrad"]:
sys.exit("GPU version of Adagrad is not supported by PyTorch.")
# specify the optimizer algorithm
opts = {
"sgd": torch.optim.SGD,
"rwsadagrad": RowWiseSparseAdagrad.RWSAdagrad,
"adagrad": torch.optim.Adagrad,
}
parameters = (
dlrm.parameters()
if ext_dist.my_size == 1
else [
{
"params": [p for emb in dlrm.emb_l for p in emb.parameters()],
"lr": args.learning_rate,
},
# TODO check this lr setup
# bottom mlp has no data parallelism
# need to check how do we deal with top mlp
{
"params": dlrm.bot_l.parameters(),
"lr": args.learning_rate,
},
{
"params": dlrm.top_l.parameters(),
"lr": args.learning_rate,
},
]
)
optimizer = opts[args.optimizer](parameters, lr=args.learning_rate)
lr_scheduler = LRPolicyScheduler(
optimizer,
args.lr_num_warmup_steps,
args.lr_decay_start_step,
args.lr_num_decay_steps,
)
### main loop ###
# training or inference
best_acc_test = 0
best_auc_test = 0
skip_upto_epoch = 0
skip_upto_batch = 0
total_time = 0
total_loss = 0
total_iter = 0
total_samp = 0
if args.mlperf_logging:
mlperf_logger.mlperf_submission_log("dlrm")
mlperf_logger.log_event(
key=mlperf_logger.constants.SEED, value=args.numpy_rand_seed
)
mlperf_logger.log_event(
key=mlperf_logger.constants.GLOBAL_BATCH_SIZE, value=args.mini_batch_size
)
# Load model is specified
if not (args.load_model == ""):
print("Loading saved model {}".format(args.load_model))
if 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(args.load_model)
else:
# NOTE: when targeting inference on single GPU,
# note that the call to .to(device) has already happened
ld_model = torch.load(
args.load_model,
map_location=torch.device("cuda")
# map_location=lambda storage, loc: storage.cuda(0)
)
else:
# when targeting inference on CPU
ld_model = torch.load(args.load_model, map_location=torch.device("cpu"))
dlrm.load_state_dict(ld_model["state_dict"])
ld_j = ld_model["iter"]
ld_k = ld_model["epoch"]
ld_nepochs = ld_model["nepochs"]
ld_nbatches = ld_model["nbatches"]
ld_nbatches_test = ld_model["nbatches_test"]
ld_train_loss = ld_model["train_loss"]
ld_total_loss = ld_model["total_loss"]
if args.mlperf_logging:
ld_gAUC_test = ld_model["test_auc"]
ld_acc_test = ld_model["test_acc"]
if not args.inference_only:
optimizer.load_state_dict(ld_model["opt_state_dict"])
best_acc_test = ld_acc_test
total_loss = ld_total_loss
skip_upto_epoch = ld_k # epochs
skip_upto_batch = ld_j # batches
else:
args.print_freq = ld_nbatches
args.test_freq = 0
print(
"Saved at: epoch = {:d}/{:d}, batch = {:d}/{:d}, ntbatch = {:d}".format(
ld_k, ld_nepochs, ld_j, ld_nbatches, ld_nbatches_test
)
)
print(
"Training state: loss = {:.6f}".format(
ld_train_loss,
)
)
if args.mlperf_logging:
print(
"Testing state: accuracy = {:3.3f} %, auc = {:.3f}".format(
ld_acc_test * 100, ld_gAUC_test
)
)
else:
print("Testing state: accuracy = {:3.3f} %".format(ld_acc_test * 100))
if args.inference_only:
# Currently only dynamic quantization with INT8 and FP16 weights are
# supported for MLPs and INT4 and INT8 weights for EmbeddingBag
# post-training quantization during the inference.
# By default we don't do the quantization: quantize_{mlp,emb}_with_bit == 32 (FP32)
assert args.quantize_mlp_with_bit in [
8,
16,
32,
], "only support 8/16/32-bit but got {}".format(args.quantize_mlp_with_bit)
assert args.quantize_emb_with_bit in [
4,
8,
32,
], "only support 4/8/32-bit but got {}".format(args.quantize_emb_with_bit)
if args.quantize_mlp_with_bit != 32:
if args.quantize_mlp_with_bit in [8]:
quantize_dtype = torch.qint8
else:
quantize_dtype = torch.float16
dlrm = torch.quantization.quantize_dynamic(
dlrm, {torch.nn.Linear}, quantize_dtype
)
if args.quantize_emb_with_bit != 32:
dlrm.quantize_embedding(args.quantize_emb_with_bit)
# print(dlrm)
print("time/loss/accuracy (if enabled):")
if args.mlperf_logging:
# LR is logged twice for now because of a compliance checker bug
mlperf_logger.log_event(
key=mlperf_logger.constants.OPT_BASE_LR, value=args.learning_rate
)
mlperf_logger.log_event(
key=mlperf_logger.constants.OPT_LR_WARMUP_STEPS,
value=args.lr_num_warmup_steps,
)
# use logging keys from the official HP table and not from the logging library
mlperf_logger.log_event(
key="sgd_opt_base_learning_rate", value=args.learning_rate
)
mlperf_logger.log_event(
key="lr_decay_start_steps", value=args.lr_decay_start_step
)
mlperf_logger.log_event(
key="sgd_opt_learning_rate_decay_steps", value=args.lr_num_decay_steps
)
mlperf_logger.log_event(key="sgd_opt_learning_rate_decay_poly_power", value=2)
tb_file = "./" + args.tensor_board_filename
writer = SummaryWriter(tb_file)
ext_dist.barrier()
with torch.autograd.profiler.profile(
args.enable_profiling, use_cuda=use_gpu, record_shapes=True
) as prof:
if not args.inference_only:
k = 0
total_time_begin = 0
while k < args.nepochs:
if args.mlperf_logging:
mlperf_logger.barrier()
mlperf_logger.log_start(
key=mlperf_logger.constants.BLOCK_START,
metadata={
mlperf_logger.constants.FIRST_EPOCH_NUM: (k + 1),
mlperf_logger.constants.EPOCH_COUNT: 1,
},
)
mlperf_logger.barrier()
mlperf_logger.log_start(
key=mlperf_logger.constants.EPOCH_START,
metadata={mlperf_logger.constants.EPOCH_NUM: (k + 1)},
)
if k < skip_upto_epoch:
continue
if args.mlperf_logging:
previous_iteration_time = None
for j, inputBatch in enumerate(train_ld):
if j == 0 and args.save_onnx:
X_onnx, lS_o_onnx, lS_i_onnx, _, _, _ = unpack_batch(inputBatch)
if j < skip_upto_batch:
continue
X, lS_o, lS_i, T, W, CBPP = unpack_batch(inputBatch)
if args.mlperf_logging:
current_time = time_wrap(use_gpu)
if previous_iteration_time:
iteration_time = current_time - previous_iteration_time
else:
iteration_time = 0
previous_iteration_time = current_time
else:
t1 = time_wrap(use_gpu)
# early exit if nbatches was set by the user and has been exceeded
if nbatches > 0 and j >= nbatches:
break
# Skip the batch if batch size not multiple of total ranks
if ext_dist.my_size > 1 and X.size(0) % ext_dist.my_size != 0:
print(
"Warning: Skiping the batch %d with size %d"
% (j, X.size(0))
)
continue
mbs = T.shape[0] # = args.mini_batch_size except maybe for last
# forward pass
Z = dlrm_wrap(
X,
lS_o,
lS_i,
use_gpu,
device,
ndevices=ndevices,
)
if ext_dist.my_size > 1:
T = T[ext_dist.get_my_slice(mbs)]
W = W[ext_dist.get_my_slice(mbs)]
# loss
E = loss_fn_wrap(Z, T, use_gpu, device)
# compute loss and accuracy
L = E.detach().cpu().numpy() # numpy array
# training accuracy is not disabled
# S = Z.detach().cpu().numpy() # numpy array
# T = T.detach().cpu().numpy() # numpy array
# # print("res: ", S)
# # print("j, train: BCE ", j, L)
# mbs = T.shape[0] # = args.mini_batch_size except maybe for last
# A = np.sum((np.round(S, 0) == T).astype(np.uint8))
with record_function("DLRM backward"):
# scaled error gradient propagation
# (where we do not accumulate gradients across mini-batches)
if (args.mlperf_logging and (j + 1) % args.mlperf_grad_accum_iter == 0) or not args.mlperf_logging:
optimizer.zero_grad()
# backward pass
E.backward()
# optimizer
if (args.mlperf_logging and (j + 1) % args.mlperf_grad_accum_iter == 0) or not args.mlperf_logging:
optimizer.step()
lr_scheduler.step()
if args.mlperf_logging:
total_time += iteration_time
else:
t2 = time_wrap(use_gpu)
total_time += t2 - t1
total_loss += L * mbs
total_iter += 1
total_samp += mbs
should_print = ((j + 1) % args.print_freq == 0) or (
j + 1 == nbatches
)
should_test = (
(args.test_freq > 0)
and (args.data_generation in ["dataset", "random"])
and (((j + 1) % args.test_freq == 0) or (j + 1 == nbatches))
)
# print time, loss and accuracy
if should_print or should_test:
gT = 1000.0 * total_time / total_iter if args.print_time else -1
total_time = 0
train_loss = total_loss / total_samp
total_loss = 0
str_run_type = (
"inference" if args.inference_only else "training"
)
wall_time = ""
if args.print_wall_time:
wall_time = " ({})".format(time.strftime("%H:%M"))
print(
"Finished {} it {}/{} of epoch {}, {:.2f} ms/it,".format(
str_run_type, j + 1, nbatches, k, gT
)
+ " loss {:.6f}".format(train_loss)
+ wall_time,
flush=True,
)
log_iter = nbatches * k + j + 1
writer.add_scalar("Train/Loss", train_loss, log_iter)
total_iter = 0
total_samp = 0
# testing
if should_test:
epoch_num_float = (j + 1) / len(train_ld) + k + 1
if args.mlperf_logging:
mlperf_logger.barrier()
mlperf_logger.log_start(
key=mlperf_logger.constants.EVAL_START,
metadata={
mlperf_logger.constants.EPOCH_NUM: epoch_num_float
},
)
# don't measure training iter time in a test iteration
if args.mlperf_logging:
previous_iteration_time = None
print(
"Testing at - {}/{} of epoch {},".format(j + 1, nbatches, k)
)
model_metrics_dict, is_best = inference(
args,
dlrm,
best_acc_test,
best_auc_test,
test_ld,
device,
use_gpu,
log_iter,
)
if (
is_best
and not (args.save_model == "")
and not args.inference_only
):
model_metrics_dict["epoch"] = k
model_metrics_dict["iter"] = j + 1
model_metrics_dict["train_loss"] = train_loss
model_metrics_dict["total_loss"] = total_loss
model_metrics_dict[
"opt_state_dict"
] = optimizer.state_dict()
print("Saving model to {}".format(args.save_model))
torch.save(model_metrics_dict, args.save_model)
if args.mlperf_logging:
mlperf_logger.barrier()
mlperf_logger.log_end(
key=mlperf_logger.constants.EVAL_STOP,
metadata={
mlperf_logger.constants.EPOCH_NUM: epoch_num_float
},
)
# Uncomment the line below to print out the total time with overhead
# print("Total test time for this group: {}" \
# .format(time_wrap(use_gpu) - accum_test_time_begin))
if (
args.mlperf_logging
and (args.mlperf_acc_threshold > 0)
and (best_acc_test > args.mlperf_acc_threshold)
):
print(
"MLPerf testing accuracy threshold "
+ str(args.mlperf_acc_threshold)
+ " reached, stop training"
)
break
if (
args.mlperf_logging
and (args.mlperf_auc_threshold > 0)
and (best_auc_test > args.mlperf_auc_threshold)
):
print(
"MLPerf testing auc threshold "
+ str(args.mlperf_auc_threshold)
+ " reached, stop training"
)
if args.mlperf_logging:
mlperf_logger.barrier()
mlperf_logger.log_end(
key=mlperf_logger.constants.RUN_STOP,
metadata={
mlperf_logger.constants.STATUS: mlperf_logger.constants.SUCCESS
},
)
break
if args.mlperf_logging:
mlperf_logger.barrier()
mlperf_logger.log_end(
key=mlperf_logger.constants.EPOCH_STOP,
metadata={mlperf_logger.constants.EPOCH_NUM: (k + 1)},
)
mlperf_logger.barrier()
mlperf_logger.log_end(
key=mlperf_logger.constants.BLOCK_STOP,
metadata={mlperf_logger.constants.FIRST_EPOCH_NUM: (k + 1)},
)
k += 1 # nepochs
if args.mlperf_logging and best_auc_test <= args.mlperf_auc_threshold:
mlperf_logger.barrier()
mlperf_logger.log_end(
key=mlperf_logger.constants.RUN_STOP,
metadata={
mlperf_logger.constants.STATUS: mlperf_logger.constants.ABORTED
},
)
else:
print("Testing for inference only")
inference(
args,
dlrm,
best_acc_test,
best_auc_test,
test_ld,
device,
use_gpu,
)
# profiling
if args.enable_profiling:
time_stamp = str(datetime.datetime.now()).replace(" ", "_")
with open("dlrm_s_pytorch" + time_stamp + "_shape.prof", "w") as prof_f:
prof_f.write(
prof.key_averages(group_by_input_shape=True).table(
sort_by="self_cpu_time_total"
)
)
with open("dlrm_s_pytorch" + time_stamp + "_total.prof", "w") as prof_f:
prof_f.write(prof.key_averages().table(sort_by="self_cpu_time_total"))
prof.export_chrome_trace("dlrm_s_pytorch" + time_stamp + ".json")
# print(prof.key_averages().table(sort_by="cpu_time_total"))
# plot compute graph
if args.plot_compute_graph:
sys.exit(
"ERROR: Please install pytorchviz package in order to use the"
+ " visualization. Then, uncomment its import above as well as"
+ " three lines below and run the code again."
)
# V = Z.mean() if args.inference_only else E
# dot = make_dot(V, params=dict(dlrm.named_parameters()))
# dot.render('dlrm_s_pytorch_graph') # write .pdf file
# test prints
if not args.inference_only and args.debug_mode:
print("updated parameters (weights and bias):")
for param in dlrm.parameters():
print(param.detach().cpu().numpy())
# export the model in onnx
if args.save_onnx:
"""
# workaround 1: tensor -> list
if torch.is_tensor(lS_i_onnx):
lS_i_onnx = [lS_i_onnx[j] for j in range(len(lS_i_onnx))]
# workaound 2: list -> tensor
lS_i_onnx = torch.stack(lS_i_onnx)
"""
# debug prints
# print("inputs", X_onnx, lS_o_onnx, lS_i_onnx)
# print("output", dlrm_wrap(X_onnx, lS_o_onnx, lS_i_onnx, use_gpu, device))
dlrm_pytorch_onnx_file = "dlrm_s_pytorch.onnx"
batch_size = X_onnx.shape[0]
print("X_onnx.shape", X_onnx.shape)
if torch.is_tensor(lS_o_onnx):
print("lS_o_onnx.shape", lS_o_onnx.shape)
else:
for oo in lS_o_onnx:
print("oo.shape", oo.shape)
if torch.is_tensor(lS_i_onnx):
print("lS_i_onnx.shape", lS_i_onnx.shape)
else:
for ii in lS_i_onnx:
print("ii.shape", ii.shape)
# name inputs and outputs
o_inputs = (
["offsets"]
if torch.is_tensor(lS_o_onnx)
else ["offsets_" + str(i) for i in range(len(lS_o_onnx))]
)
i_inputs = (
["indices"]
if torch.is_tensor(lS_i_onnx)
else ["indices_" + str(i) for i in range(len(lS_i_onnx))]
)
all_inputs = ["dense_x"] + o_inputs + i_inputs
# debug prints
print("inputs", all_inputs)
# create dynamic_axis dictionaries
do_inputs = (
[{"offsets": {1: "batch_size"}}]
if torch.is_tensor(lS_o_onnx)
else [
{"offsets_" + str(i): {0: "batch_size"}} for i in range(len(lS_o_onnx))
]
)
di_inputs = (
[{"indices": {1: "batch_size"}}]
if torch.is_tensor(lS_i_onnx)
else [
{"indices_" + str(i): {0: "batch_size"}} for i in range(len(lS_i_onnx))
]
)
dynamic_axes = {"dense_x": {0: "batch_size"}, "pred": {0: "batch_size"}}
for do in do_inputs:
dynamic_axes.update(do)
for di in di_inputs:
dynamic_axes.update(di)
# debug prints
print(dynamic_axes)
# export model
torch.onnx.export(
dlrm,
(X_onnx, lS_o_onnx, lS_i_onnx),
dlrm_pytorch_onnx_file,
verbose=True,
opset_version=11,
input_names=all_inputs,
output_names=["pred"],
dynamic_axes=dynamic_axes,
)
# recover the model back
dlrm_pytorch_onnx = onnx.load("dlrm_s_pytorch.onnx")
# check the onnx model
onnx.checker.check_model(dlrm_pytorch_onnx)
total_time_end = time_wrap(use_gpu)
if __name__ == "__main__":
run()
PyTorch/Recommendation/dlrm/extend_distributed.py 0 → 100644
View file @ c43a53e4
# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
import builtins
import os
import sys
import torch
import torch.distributed as dist
from torch.autograd import Function
from torch.autograd.profiler import record_function
from torch.nn.parallel import DistributedDataParallel as DDP
try:
import torch_ccl
except ImportError as e:
# print(e)
torch_ccl = False
try:
import torch_ucc
except ImportError as e:
torch_ucc = False
my_rank = -1
my_size = -1
my_local_rank = -1
my_local_size = -1
alltoall_supported = False
a2a_impl = os.environ.get("DLRM_ALLTOALL_IMPL", "")
myreq = None
def env2int(env_list, default=-1):
for e in env_list:
val = int(os.environ.get(e, -1))
if val >= 0:
return val
return default
def get_my_slice(n):
k, m = divmod(n, my_size)
return slice(
my_rank * k + min(my_rank, m), (my_rank + 1) * k + min(my_rank + 1, m), 1
)
def get_split_lengths(n):
k, m = divmod(n, my_size)
if m == 0:
splits = None
my_len = k
else:
splits = [(k + 1) if i < m else k for i in range(my_size)]
my_len = splits[my_rank]
return (my_len, splits)
def init_distributed(rank=-1, local_rank=-1, size=-1, use_gpu=False, backend=""):
global myreq
global my_rank
global my_size
global my_local_rank
global my_local_size
global a2a_impl
global alltoall_supported
# guess MPI ranks from env (works for IMPI, OMPI and MVAPICH2)
num_mpi_ranks = env2int(
["PMI_SIZE", "OMPI_COMM_WORLD_SIZE", "MV2_COMM_WORLD_SIZE", "WORLD_SIZE"]
)
if backend == "" and num_mpi_ranks > 1:
if torch_ccl and env2int(["CCL_WORKER_COUNT"]) > 0:
backend = "ccl"
elif use_gpu and dist.is_nccl_available():
backend = "nccl"
elif dist.is_mpi_available():
backend = "mpi"
else:
print(
"WARNING: MPI multi-process launch detected but PyTorch MPI backend not available."
)
backend = "gloo"
if backend != "":
# guess Rank and size
if rank == -1:
rank = env2int(
["PMI_RANK", "OMPI_COMM_WORLD_RANK", "MV2_COMM_WORLD_RANK", "RANK"], 0
)
if size == -1:
size = env2int(
[
"PMI_SIZE",
"OMPI_COMM_WORLD_SIZE",
"MV2_COMM_WORLD_SIZE",
"WORLD_SIZE",
],
1,
)
if not os.environ.get("RANK", None) and rank != -1:
os.environ["RANK"] = str(rank)
if not os.environ.get("WORLD_SIZE", None) and size != -1:
os.environ["WORLD_SIZE"] = str(size)
if not os.environ.get("MASTER_PORT", None):
os.environ["MASTER_PORT"] = "29500"
if not os.environ.get("MASTER_ADDR", None):
local_size = env2int(
[
"MPI_LOCALNRANKS",
"OMPI_COMM_WORLD_LOCAL_SIZE",
"MV2_COMM_WORLD_LOCAL_SIZE",
],
1,
)
if local_size != size and backend != "mpi":
print(
"Warning: Looks like distributed multinode run but MASTER_ADDR env not set, using '127.0.0.1' as default"
)
print(
"If this run hangs, try exporting rank 0's hostname as MASTER_ADDR"
)
os.environ["MASTER_ADDR"] = "127.0.0.1"
if size > 1:
if local_rank == -1:
my_local_rank = env2int(
[
"MPI_LOCALRANKID",
"OMPI_COMM_WORLD_LOCAL_RANK",
"MV2_COMM_WORLD_LOCAL_RANK",
"LOCAL_RANK",
],
0,
)
else:
my_local_rank = local_rank
my_local_size = env2int(
[
"MPI_LOCALNRANKS",
"OMPI_COMM_WORLD_LOCAL_SIZE",
"MV2_COMM_WORLD_LOCAL_SIZE",
],
1,
)
if use_gpu:
if my_local_size > torch.cuda.device_count():
print(
"Not sufficient GPUs available... local_size = %d, ngpus = %d"
% (my_local_size, torch.cuda.device_count())
)
sys.exit(1)
torch.cuda.set_device(my_local_rank)
dist.init_process_group(backend, rank=rank, world_size=size)
my_rank = dist.get_rank()
my_size = dist.get_world_size()
if my_rank == 0:
print("Running on %d ranks using %s backend" % (my_size, backend))
if hasattr(dist, "all_to_all_single"):
try:
t = torch.zeros([4])
if use_gpu:
t = t.cuda()
dist.all_to_all_single(t, t)
alltoall_supported = True
except RuntimeError as err:
print("fail to enable all_to_all_single primitive: %s" % err)
if a2a_impl == "alltoall" and alltoall_supported == False:
print(
"Requested DLRM_ALLTOALL_IMPL=%s but backend %s does not support it, use scatter/gather based alltoall"
% (a2a_impl, backend)
)
a2a_impl = "scatter"
if a2a_impl != "":
print("Using DLRM_ALLTOALL_IMPL=%s" % a2a_impl)
else:
my_rank = 0
my_size = 1
my_local_rank = 0
my_local_size = 1
print_all(
"world size: %d, current rank: %d, local rank: %d"
% (my_size, my_rank, my_local_rank)
)
myreq = Request()
class Request(object):
def __init__(self):
self.req = None
self.tensor = None
self.WaitFunction = All2All_Scatter_Wait
def wait(self):
ret = self.WaitFunction.apply(*self.tensor)
self.req = None
self.tensor = None
return ret
class All2All_ScatterList_Req(Function):
@staticmethod
def forward(ctx, a2a_info, *inputs):
global myreq
batch_split_lengths = (
a2a_info.global_batch_partition_slices
if a2a_info.global_batch_partition_slices
else a2a_info.local_batch_num
)
table_split_lengths = (
a2a_info.global_table_wise_parition_slices
if a2a_info.global_table_wise_parition_slices
else [a2a_info.local_table_num] * my_size
)
gather_list = []
req_list = []
for i in range(my_size):
for j in range(table_split_lengths[i]):
out_tensor = inputs[0].new_empty(
[a2a_info.local_batch_num, a2a_info.emb_dim]
)
scatter_list = (
list(inputs[j].split(batch_split_lengths, dim=0))
if i == my_rank
else []
)
req = dist.scatter(out_tensor, scatter_list, src=i, async_op=True)
gather_list.append(out_tensor)
req_list.append(req)
myreq.req = req_list
myreq.tensor = tuple(gather_list)
myreq.a2a_info = a2a_info
return myreq.tensor
@staticmethod
def backward(ctx, *grad_output):
global myreq
for r in myreq.req:
r.wait()
myreq.req = None
grad_inputs = myreq.tensor
myreq.tensor = None
return (None, *grad_inputs)
class All2All_ScatterList_Wait(Function):
@staticmethod
def forward(ctx, *output):
global myreq
ctx.a2a_info = myreq.a2a_info
for r in myreq.req:
r.wait()
myreq.req = None
myreq.tensor = None
return output
@staticmethod
def backward(ctx, *grad_output):
global myreq
a2a_info = ctx.a2a_info
grad_output = [t.contiguous() for t in grad_output]
batch_split_lengths = (
a2a_info.global_batch_partition_slices
if a2a_info.global_batch_partition_slices
else [a2a_info.local_batch_num] * my_size
)
per_rank_table_splits = (
a2a_info.global_table_wise_parition_slices
if a2a_info.global_table_wise_parition_slices
else [a2a_info.local_table_num] * my_size
)
grad_inputs = [
grad_output[0].new_empty([ctx.a2a_info.batch_size, ctx.a2a_info.emb_dim])
for _ in range(a2a_info.local_table_num)
]
req_list = []
ind = 0
for i in range(my_size):
for j in range(per_rank_table_splits[i]):
gather_list = (
list(grad_inputs[j].split(batch_split_lengths, dim=0))
if i == my_rank
else None
)
req = dist.gather(grad_output[ind], gather_list, dst=i, async_op=True)
req_list.append(req)
ind += 1
myreq.req = req_list
myreq.tensor = grad_inputs
return tuple(grad_output)
class All2All_Scatter_Req(Function):
@staticmethod
def forward(ctx, a2a_info, *inputs):
global myreq
batch_split_lengths = (
a2a_info.global_batch_partition_slices
if a2a_info.global_batch_partition_slices
else a2a_info.local_batch_num
)
table_split_lengths = (
a2a_info.global_table_wise_parition_slices
if a2a_info.global_table_wise_parition_slices
else [a2a_info.local_table_num] * my_size
)
input = torch.cat(inputs, dim=1)
scatter_list = list(input.split(batch_split_lengths, dim=0))
gather_list = []
req_list = []
for i in range(my_size):
out_tensor = input.new_empty(
[a2a_info.local_batch_num, table_split_lengths[i] * a2a_info.emb_dim]
)
req = dist.scatter(
out_tensor, scatter_list if i == my_rank else [], src=i, async_op=True
)
gather_list.append(out_tensor)
req_list.append(req)
myreq.req = req_list
myreq.tensor = tuple(gather_list)
myreq.a2a_info = a2a_info
ctx.a2a_info = a2a_info
return myreq.tensor
@staticmethod
def backward(ctx, *grad_output):
global myreq
for r in myreq.req:
r.wait()
myreq.req = None
grad_input = myreq.tensor
grad_inputs = grad_input.split(ctx.a2a_info.emb_dim, dim=1)
myreq.tensor = None
return (None, *grad_inputs)
class All2All_Scatter_Wait(Function):
@staticmethod
def forward(ctx, *output):
global myreq
ctx.a2a_info = myreq.a2a_info
for r in myreq.req:
r.wait()
myreq.req = None
myreq.tensor = None
return output
@staticmethod
def backward(ctx, *grad_output):
global myreq
assert len(grad_output) == my_size
scatter_list = [t.contiguous() for t in grad_output]
a2a_info = ctx.a2a_info
batch_split_lengths = (
a2a_info.global_batch_partition_slices
if a2a_info.global_batch_partition_slices
else a2a_info.local_batch_num
)
table_split_lengths = (
a2a_info.global_table_wise_parition_slices
if a2a_info.global_table_wise_parition_slices
else [a2a_info.local_table_num] * my_size
)
grad_input = grad_output[0].new_empty(
[a2a_info.batch_size, a2a_info.emb_dim * a2a_info.local_table_num]
)
gather_list = list(grad_input.split(batch_split_lengths, dim=0))
req_list = []
for i in range(my_size):
req = dist.gather(
scatter_list[i],
gather_list if i == my_rank else [],
dst=i,
async_op=True,
)
req_list.append(req)
myreq.req = req_list
myreq.tensor = grad_input
return grad_output
class All2All_Req(Function):
@staticmethod
def forward(ctx, a2a_info, *inputs):
global myreq
with record_function("DLRM alltoall_req_fwd_single"):
batch_split_lengths = a2a_info.global_batch_partition_slices
if batch_split_lengths:
batch_split_lengths = [
m * a2a_info.emb_dim * a2a_info.local_table_num
for m in batch_split_lengths
]
table_split_lengths = a2a_info.global_table_wise_parition_slices
if table_split_lengths:
table_split_lengths = [
a2a_info.local_batch_num * e * a2a_info.emb_dim
for e in table_split_lengths
]
input = torch.cat(inputs, dim=1).view([-1])
output = input.new_empty(
[
a2a_info.global_table_num
* a2a_info.local_batch_num
* a2a_info.emb_dim
]
)
req = dist.all_to_all_single(
output, input, table_split_lengths, batch_split_lengths, async_op=True
)
myreq.req = req
myreq.tensor = []
myreq.tensor.append(output)
myreq.tensor = tuple(myreq.tensor)
a2a_info.batch_split_lengths = batch_split_lengths
a2a_info.table_split_lengths = table_split_lengths
myreq.a2a_info = a2a_info
ctx.a2a_info = a2a_info
return myreq.tensor
@staticmethod
def backward(ctx, *grad_output):
global myreq
with record_function("DLRM alltoall_req_bwd_single"):
a2a_info = ctx.a2a_info
myreq.req.wait()
myreq.req = None
grad_input = myreq.tensor
grad_inputs = grad_input.view([a2a_info.batch_size, -1]).split(
a2a_info.emb_dim, dim=1
)
grad_inputs = [gin.contiguous() for gin in grad_inputs]
myreq.tensor = None
return (None, *grad_inputs)
class All2All_Wait(Function):
@staticmethod
def forward(ctx, *output):
global myreq
with record_function("DLRM alltoall_wait_fwd_single"):
a2a_info = myreq.a2a_info
ctx.a2a_info = a2a_info
myreq.req.wait()
myreq.req = None
myreq.tensor = None
table_split_lengths = (
a2a_info.table_split_lengths
if a2a_info.table_split_lengths
else a2a_info.local_table_num
* a2a_info.local_batch_num
* a2a_info.emb_dim
)
outputs = output[0].split(table_split_lengths)
outputs = tuple(
[out.view([a2a_info.local_batch_num, -1]) for out in outputs]
)
return outputs
@staticmethod
def backward(ctx, *grad_outputs):
global myreq
with record_function("DLRM alltoall_wait_bwd_single"):
a2a_info = ctx.a2a_info
grad_outputs = [gout.contiguous().view([-1]) for gout in grad_outputs]
grad_output = torch.cat(grad_outputs)
grad_input = grad_output.new_empty(
[a2a_info.batch_size * a2a_info.local_table_num * a2a_info.emb_dim]
)
req = dist.all_to_all_single(
grad_input,
grad_output,
a2a_info.batch_split_lengths,
a2a_info.table_split_lengths,
async_op=True,
)
myreq.req = req
myreq.tensor = grad_input
return (grad_output,)
class AllGather(Function):
@staticmethod
def forward(ctx, input, global_lengths, dim=0):
if not isinstance(global_lengths, (list, tuple)):
global_lengths = [global_lengths] * my_size
assert len(global_lengths) == my_size
assert global_lengths[my_rank] == input.size(dim)
local_start = sum(global_lengths[:my_rank])
output_size = list(input.size())
ctx.dim = dim
ctx.local_start = local_start
ctx.local_length = global_lengths[my_rank]
input = input.contiguous()
if dim == 0:
out_len = sum(global_lengths)
output_size[dim] = out_len
output = input.new_empty(output_size)
gather_list = list(output.split(global_lengths, dim=0))
else:
gather_list = [torch.empty_like(input) for _ in range(my_size)]
gather_list = []
for length in global_lengths:
output_size[dim] = length
gather_list.append(input.new_empty(output_size))
dist.all_gather(gather_list, input)
if dim != 0:
output = torch.cat(gather_list, dim=dim)
return output
@staticmethod
def backward(ctx, grad_output):
# print("Inside All2AllBackward")
dim = ctx.dim
start = ctx.local_start
length = ctx.local_length
grad_input = grad_output.narrow(dim, start, length)
return (grad_input, None, None)
class All2AllInfo(object):
pass
def alltoall(inputs, per_rank_table_splits):
global myreq
batch_size, emb_dim = inputs[0].size()
a2a_info = All2AllInfo()
a2a_info.local_table_num = len(inputs)
a2a_info.global_table_wise_parition_slices = per_rank_table_splits
(
a2a_info.local_batch_num,
a2a_info.global_batch_partition_slices,
) = get_split_lengths(batch_size)
a2a_info.emb_dim = emb_dim
a2a_info.batch_size = batch_size
a2a_info.global_table_num = (
sum(per_rank_table_splits)
if per_rank_table_splits
else a2a_info.local_table_num * my_size
)
if a2a_impl == "" and alltoall_supported or a2a_impl == "alltoall":
# print("Using All2All_Req")
output = All2All_Req.apply(a2a_info, *inputs)
myreq.WaitFunction = All2All_Wait
elif a2a_impl == "" or a2a_impl == "scatter":
# print("Using All2All_Scatter_Req")
output = All2All_Scatter_Req.apply(a2a_info, *inputs)
myreq.WaitFunction = All2All_Scatter_Wait
elif a2a_impl == "scatter_list":
# print("Using All2All_ScatterList_Req")
output = All2All_ScatterList_Req.apply(a2a_info, *inputs)
myreq.WaitFunction = All2All_ScatterList_Wait
else:
print(
"Unknown value set for DLRM_ALLTOALL_IMPL (%s), "
"please use one of [alltoall, scatter, scatter_list]" % a2a_impl
)
return myreq
def all_gather(input, lengths, dim=0):
if not lengths:
lengths = [input.size(0)] * my_size
return AllGather.apply(input, lengths, dim)
def barrier():
if my_size > 1:
dist.barrier()
# Override builtin print function to print only from rank 0
orig_print = builtins.print
def rank0_print(*args, **kwargs):
if my_rank <= 0 or kwargs.get("print_all", False):
orig_print(*args, **kwargs)
builtins.print = rank0_print
# Allow printing from all rank with explicit print_all
def print_all(*args, **kwargs):
orig_print(*args, **kwargs)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment