forked from jiuyuan/InfiniTensor
1.4 KiB
1.4 KiB
分布式脚本
英伟达平台运行方式
1. 运行pytorch模型并生成输入和标准输出，可选择导出onnx
使用 --export_onnx 设置导出onnx的目录，默认为当前路径 ./，不使用这个flag则只进行计算和生成输入输出。
python run_pytorch.py --model gpt2 --batch_size 1 --length 1 --export_onnx ./
会在当前目录下生成输入输出文件test_inputs.npy 和 test_results.npy，目前只支持单一输入输出。
2. 运行InfiniTensor分布式脚本
python cuda_launch.py --model "/XXX/XXX.onnx" --nproc_per_node 4
寒武纪平台运行方式
将上述运行脚本 run_pytorch.py 以及 cuda_launch.py 针对寒武纪平台做了相应的适配，具体见 run_pytorch_mlu.py 以及 bang_launch.py。
1. 运行pytorch模型并生成输入和标准输出，可选择导出onnx
使用 --export_onnx 设置导出onnx的目录，默认为当前路径 ./，不使用这个flag则只进行计算和生成输入输出。
python run_pytorch_mlu.py --model gpt2 --batch_size 1 --length 1 --export_onnx ./
会在当前目录下生成输入输出文件test_inputs.npy 和 test_results.npy，目前只支持单一输入输出。
2. 运行InfiniTensor分布式脚本
python bang_launch.py --model "/XXX/XXX.onnx" --nproc_per_node 4