From 6d6b6b54a544f69fb201a9aab7097358f6a440ff Mon Sep 17 00:00:00 2001 From: Hariharan Seshadri Date: Tue, 7 Jul 2020 21:09:37 -0700 Subject: [PATCH] Support binding a graph output to a specific device via the Python binding (#4439) --- onnxruntime/core/session/IOBinding.cc | 11 ++- onnxruntime/core/session/IOBinding.h | 24 +++-- onnxruntime/core/session/inference_session.cc | 7 +- onnxruntime/core/session/inference_session.h | 5 + .../python/onnxruntime_pybind_state.cc | 70 +++++++------ onnxruntime/python/session.py | 21 +++- .../test/framework/inference_session_test.cc | 2 +- .../onnxruntime_test_python_iobinding.py | 98 ++++++++++++++++--- orttraining/orttraining/python/ort_trainer.py | 5 +- tools/ci_build/build.py | 16 +++ .../linux/docker/scripts/install_deps.sh | 2 +- 11 files changed, 194 insertions(+), 67 deletions(-) diff --git a/onnxruntime/core/session/IOBinding.cc b/onnxruntime/core/session/IOBinding.cc index 053c221ff1..c77b37507d 100644 --- a/onnxruntime/core/session/IOBinding.cc +++ b/onnxruntime/core/session/IOBinding.cc @@ -78,7 +78,16 @@ common::Status IOBinding::SynchronizeOutputs() { return Status::OK(); } -common::Status IOBinding::BindOutput(const std::string& name, const OrtValue& ml_value, OrtDevice device) { +common::Status IOBinding::BindOutput(const std::string& name, const OrtValue& ml_value) { + // device value is ignored when ml_value is pre-allocated + return BindOutputImpl(name, ml_value, {}); +} + +common::Status IOBinding::BindOutput(const std::string& name, OrtDevice device) { + return BindOutputImpl(name, {}, device); +} + +common::Status IOBinding::BindOutputImpl(const std::string& name, const OrtValue& ml_value, OrtDevice device) { auto rc = Contains(output_names_, name); if (rc.first) { outputs_[rc.second] = ml_value; diff --git a/onnxruntime/core/session/IOBinding.h b/onnxruntime/core/session/IOBinding.h index 150db507cd..62953f0178 100644 --- a/onnxruntime/core/session/IOBinding.h +++ b/onnxruntime/core/session/IOBinding.h @@ -60,13 +60,16 @@ class IOBinding { common::Status SynchronizeOutputs(); /** - * Bind an output name to a provided OrtValue. - * If the output is pre-allocated, the value in 'device' is not used. - * If the output is not pre-allocated, information on the device it should be allocated on can be provided. - * - * @param device Device to allocate on if not pre-allocated. Default is CPU. + * Bind an output name to a provided pre-allocated OrtValue. */ - common::Status BindOutput(const std::string& name, const OrtValue& ml_value, OrtDevice device = {}); + common::Status BindOutput(const std::string& name, const OrtValue& ml_value); + + /** + * Bind an output name to a device. + * + * @param device Device to allocate the output on. Default is CPU. + */ + common::Status BindOutput(const std::string& name, OrtDevice device = {}); /** * This simply collects the outputs obtained after calling Run() inside the @param outputs. @@ -74,9 +77,6 @@ class IOBinding { const std::vector& GetOutputNames() const; std::vector& GetOutputs(); - // device info for outputs that are not pre-allocated - const std::vector& GetOutputsDeviceInfo() const; - const std::vector& GetInputNames() const; const std::vector& GetInputs() const; @@ -105,5 +105,11 @@ class IOBinding { std::vector outputs_device_info_; ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(IOBinding); + + // device info for all outputs. only used by InferenceSession if the output is not pre-allocated. + const std::vector& GetOutputsDeviceInfo() const; + + // The implementation for the BindOutput() overloads + common::Status BindOutputImpl(const std::string& name, const OrtValue& ml_value, OrtDevice device); }; } // namespace onnxruntime diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 702d2cd0ba..aef1e7a551 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -991,6 +991,10 @@ const SessionOptions& InferenceSession::GetSessionOptions() const { return session_options_; } +const DataTransferManager& InferenceSession::GetDataTransferManager() const { + return data_transfer_mgr_; +} + common::Status InferenceSession::CheckShapes(const std::string& input_name, const TensorShape& input_shape, const TensorShape& expected_shape) const { auto input_shape_sz = input_shape.NumDimensions(); @@ -1546,9 +1550,8 @@ common::Status InferenceSession::WaitForNotification(Notification* p_executor_do return Status::OK(); } -SessionIOBinding::SessionIOBinding(InferenceSession* session) { +SessionIOBinding::SessionIOBinding(InferenceSession* session) : sess_(session) { ORT_ENFORCE(session->NewIOBinding(&binding_).IsOK()); - sess_ = session; } InferenceSession* SessionIOBinding::GetInferenceSession() { diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index c34d54c5b7..16bd3b653a 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -319,6 +319,11 @@ class InferenceSession { */ const SessionOptions& GetSessionOptions() const; + /* + * Get the DataTransferManager associated with this session + */ + const DataTransferManager& GetDataTransferManager() const; + /** * Start profiling on this inference session. This simply turns on profiling events to be * recorded. A corresponding EndProfiling has to follow to write profiling data to a file. diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 2b0d7cbc72..d3d983a37b 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -185,7 +185,7 @@ using namespace onnxruntime; using namespace onnxruntime::logging; template -void AddNonTensor(const OrtValue& val, std::vector& pyobjs) { +void AddNonTensor(const OrtValue& val, std::vector& pyobjs, const DataTransferManager* /*data_transfer_manager*/) { pyobjs.push_back(py::cast(val.Get())); } @@ -242,47 +242,47 @@ static const char* GetDeviceName(const OrtDevice& device) { } template <> -void AddNonTensor(const OrtValue& val, std::vector& pyobjs) { +void AddNonTensor(const OrtValue& val, std::vector& pyobjs, const DataTransferManager* data_transfer_manager) { const auto& seq_tensors = val.Get(); py::list py_list; for (const auto& rtensor : seq_tensors) { py::object obj; - GetPyObjFromTensor(rtensor, obj, nullptr); + GetPyObjFromTensor(rtensor, obj, data_transfer_manager); py_list.append(obj); } pyobjs.push_back(py_list); } -void AddNonTensorAsPyObj(const OrtValue& val, std::vector& pyobjs) { +void AddNonTensorAsPyObj(const OrtValue& val, std::vector& pyobjs, const DataTransferManager* data_transfer_manager) { // Should be in sync with core/framework/datatypes.h auto val_type = val.Type(); if (val_type->IsTensorSequenceType()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else { utils::ContainerChecker c_checker(val_type); if (c_checker.IsMap()) { if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsMapOf()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } } else { if (c_checker.IsSequenceOf>()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else if (c_checker.IsSequenceOf>()) { - AddNonTensor(val, pyobjs); + AddNonTensor(val, pyobjs, data_transfer_manager); } else { throw std::runtime_error("Output is a non-tensor type which is not supported."); } @@ -290,10 +290,10 @@ void AddNonTensorAsPyObj(const OrtValue& val, std::vector& pyobjs) { } } -void AddTensorAsPyObj(const OrtValue& val, std::vector& pyobjs) { +void AddTensorAsPyObj(const OrtValue& val, std::vector& pyobjs, const DataTransferManager* data_transfer_manager) { const Tensor& rtensor = val.Get(); py::object obj; - GetPyObjFromTensor(rtensor, obj); + GetPyObjFromTensor(rtensor, obj, data_transfer_manager); pyobjs.push_back(obj); } @@ -460,13 +460,13 @@ void addGlobalMethods(py::module& m, const Environment& env) { onnxruntime::CreateExecutionProviderFactory_MIGraphX(0) #endif #ifdef USE_VITISAI - onnxruntime::CreateExecutionProviderFactory_VitisAI("DPU", 0), + onnxruntime::CreateExecutionProviderFactory_VitisAI("DPU", 0), #endif #ifdef USE_ACL onnxruntime::CreateExecutionProviderFactory_ACL(0) #endif #ifdef USE_ARMNN - onnxruntime::CreateExecutionProviderFactory_ArmNN(0) + onnxruntime::CreateExecutionProviderFactory_ArmNN(0) #endif }; @@ -642,7 +642,7 @@ void addObjectMethods(py::module& m, Environment& env) { throw std::runtime_error("Either failed to get model inputs from the session object or the input def list was null"); } - //Check if input is sequence of tensors + //Check if input is sequence of tensors onnx::TypeProto type_proto; const auto& def_list = *px.second; auto ret_it = std::find_if(std::begin(def_list), std::end(def_list), @@ -654,7 +654,7 @@ void addObjectMethods(py::module& m, Environment& env) { if (!temp) { throw std::runtime_error("Corresponding type_proto is null"); } else { - type_proto = *temp; + type_proto = *temp; } if (type_proto.has_sequence_type()) { throw std::runtime_error("Cannot bind input to sequence of tensors"); @@ -681,8 +681,8 @@ void addObjectMethods(py::module& m, Environment& env) { OrtValue mlvalue; mlvalue.Init(p_tensor.release(), - DataTypeImpl::GetType(), - DataTypeImpl::GetType()->GetDeleteFunc()); + DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); auto status = io_binding->Get()->BindInput(name, mlvalue); if (!status.IsOK()) throw std::runtime_error("Error when bind input: " + status.ErrorMessage()); @@ -699,16 +699,15 @@ void addObjectMethods(py::module& m, Environment& env) { std::unique_ptr p_tensor = onnxruntime::make_unique(NumpyTypeToOnnxRuntimeType(type_num), shape, (void*)data_ptr, info); OrtValue mlvalue; mlvalue.Init(p_tensor.release(), - DataTypeImpl::GetType(), - DataTypeImpl::GetType()->GetDeleteFunc()); - + DataTypeImpl::GetType(), + DataTypeImpl::GetType()->GetDeleteFunc()); + auto status = io_binding->Get()->BindOutput(name, mlvalue); if (!status.IsOK()) throw std::runtime_error("Error when bind output: " + status.ErrorMessage()); }) - .def("bind_output", [](SessionIOBinding* io_binding, const std::string& name) -> void { - OrtValue mlvalue; - auto status = io_binding->Get()->BindOutput(name, mlvalue); + .def("bind_output", [](SessionIOBinding* io_binding, const std::string& name, const OrtDevice& device) -> void { + auto status = io_binding->Get()->BindOutput(name, device); if (!status.IsOK()) throw std::runtime_error("Error when bind output: " + status.ErrorMessage()); }) @@ -718,15 +717,15 @@ void addObjectMethods(py::module& m, Environment& env) { .def("clear_binding_outputs", [](SessionIOBinding* io_binding) -> void { io_binding->Get()->ClearOutputs(); }) - .def("get_outputs", [](SessionIOBinding* io_binding) -> std::vector { + .def("copy_outputs_to_cpu", [](SessionIOBinding* io_binding) -> std::vector { const std::vector& outputs = io_binding->Get()->GetOutputs(); std::vector rfetch; rfetch.reserve(outputs.size()); for (const auto& _ : outputs) { if (_.IsTensor()) { - AddTensorAsPyObj(_, rfetch); + AddTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager()); } else { - AddNonTensorAsPyObj(_, rfetch); + AddNonTensorAsPyObj(_, rfetch, &io_binding->GetInferenceSession()->GetDataTransferManager()); } } return rfetch; @@ -802,8 +801,7 @@ Applies to session load, initialization, etc. Default is 0.)pbdoc") }, R"pbdoc(Graph optimization level for this session.)pbdoc") .def_readwrite("use_deterministic_compute", &SessionOptions::use_deterministic_compute, - R"pbdoc(Whether to use deterministic compute. Default is false.)pbdoc") -; + R"pbdoc(Whether to use deterministic compute. Default is false.)pbdoc"); py::class_(m, "RunOptions", R"pbdoc(Configuration information for a single Run.)pbdoc") .def(py::init()) @@ -955,9 +953,9 @@ including arg name, arg type (contains both type and shape).)pbdoc") rfetch.reserve(fetches.size()); for (auto _ : fetches) { if (_.IsTensor()) { - AddTensorAsPyObj(_, rfetch); + AddTensorAsPyObj(_, rfetch, nullptr); } else { - AddNonTensorAsPyObj(_, rfetch); + AddNonTensorAsPyObj(_, rfetch, nullptr); } } return rfetch; diff --git a/onnxruntime/python/session.py b/onnxruntime/python/session.py index dfc4cc4133..b4b620a4e5 100644 --- a/onnxruntime/python/session.py +++ b/onnxruntime/python/session.py @@ -228,17 +228,28 @@ class IOBinding: :param shape: output shape :param buffer_ptr: memory pointer to output data ''' - if device_type == 'cpu': - self._iobinding.bind_output(name) + + # Follow the `if` path when the user has not provided any pre-allocated buffer but still + # would like to bind an output to a specific device (e.g. cuda). + # Pre-allocating an output buffer may not be an option for the user as : + # (1) They may not want to use a custom allocator specific to the device they want to bind the output to, + # in which case ORT will allocate the memory for the user + # (2) The output has a dynamic shape and hence the size of the buffer may not be fixed across runs + if buffer_ptr is None: + self._iobinding.bind_output(name, + C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(), + device_id)) else: + if element_type is None or shape is None: + raise ValueError("`element_type` and `shape` are to be provided if pre-allocated memory is provided") self._iobinding.bind_output(name, C.OrtDevice(get_ort_device_type(device_type), C.OrtDevice.default_memory(), device_id), element_type, shape, buffer_ptr) - def get_outputs(self): - '''Obtain outputs.''' - return self._iobinding.get_outputs() + def copy_outputs_to_cpu(self): + '''Copy output contents to CPU (if on another device). No-op if already on the CPU.''' + return self._iobinding.copy_outputs_to_cpu() def clear_binding_inputs(self): self._iobinding.clear_binding_inputs() diff --git a/onnxruntime/test/framework/inference_session_test.cc b/onnxruntime/test/framework/inference_session_test.cc index dbd008ea10..4ec5fd2fe9 100644 --- a/onnxruntime/test/framework/inference_session_test.cc +++ b/onnxruntime/test/framework/inference_session_test.cc @@ -308,7 +308,7 @@ void RunModelWithBindingMatMul(InferenceSession& session_object, if (output_device) { // output should be allocated on specified device (if not preallocated here) - io_binding->BindOutput("Y", output_ml_value, *output_device); + io_binding->BindOutput("Y", *output_device); } else { io_binding->BindOutput("Y", output_ml_value); } diff --git a/onnxruntime/test/python/onnxruntime_test_python_iobinding.py b/onnxruntime/test/python/onnxruntime_test_python_iobinding.py index 491c63926c..27f5c9f11c 100644 --- a/onnxruntime/test/python/onnxruntime_test_python_iobinding.py +++ b/onnxruntime/test/python/onnxruntime_test_python_iobinding.py @@ -33,7 +33,8 @@ class TestIOBinding(unittest.TestCase): # Run Pytorch touch_input = torch.randn(batch_size, channels, sample_dim, sample_dim).to(device) - torch_output = model(touch_input).cpu().detach().numpy() + torch_output = model(touch_input) + torch_output_vals = torch_output.cpu().detach().numpy() # Run ORT input_names = [ "input" ] @@ -42,31 +43,106 @@ class TestIOBinding(unittest.TestCase): input_names=input_names, output_names=output_names, dynamic_axes={"input":{0:"batch_size"}, "output":{0:"batch_size"}}) - return touch_input, torch_output + return touch_input, torch_output, torch_output_vals def test_bind_input_only(self): - x, torch_output = self.create_model_and_input() + torch_input, torch_output, torch_output_vals = self.create_model_and_input() session = onnxruntime.InferenceSession('model.onnx') io_binding = session.io_binding() - io_binding.bind_input('input', x.device.type, 0, np.float32, list(x.size()), x.data_ptr()) + + # Bind input to CUDA + io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr()) + + # Bind output to CPU io_binding.bind_output('output') + + # Invoke Run session.run_with_iobinding(io_binding) - ort_output = io_binding.get_outputs()[0] - - self.assertTrue(np.array_equal(torch_output, ort_output)) + + # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) + ort_output = io_binding.copy_outputs_to_cpu()[0] + + # Validate results + self.assertTrue(np.array_equal(torch_output_vals, ort_output)) def test_bind_input_to_cpu_arr(self): - torch_input, torch_output = self.create_model_and_input() + torch_input, torch_output, torch_output_vals = self.create_model_and_input() session = onnxruntime.InferenceSession('model.onnx') io_binding = session.io_binding() + + # Bind Numpy object (input) to CUDA io_binding.bind_cpu_input('input', torch_input.cpu().detach().numpy()) + + # Bind output to CPU io_binding.bind_output('output') + + # Invoke Run session.run_with_iobinding(io_binding) - ort_output = io_binding.get_outputs()[0] - - self.assertTrue(np.array_equal(torch_output, ort_output)) + + # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) + ort_output = io_binding.copy_outputs_to_cpu()[0] + + # Validate results + self.assertTrue(np.array_equal(torch_output_vals, ort_output)) + + def test_bind_input_and_preallocated_output(self): + torch_input, torch_output, torch_output_vals = self.create_model_and_input() + + session = onnxruntime.InferenceSession('model.onnx') + io_binding = session.io_binding() + + # Bind input to CUDA + io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr()) + + # Bind output to CUDA + ort_output = torch.empty(list(torch_output.size())).cuda() + io_binding.bind_output('output', torch_output.device.type, 0, np.float32, list(torch_output.size()), ort_output.data_ptr()) + + # Invoke Run + session.run_with_iobinding(io_binding) + + # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) + ort_output_vals = io_binding.copy_outputs_to_cpu()[0] + # Validate results + self.assertTrue(np.array_equal(torch_output_vals, ort_output_vals)) + + # Validate if ORT actually wrote to pre-allocated buffer by copying the Torch allocated buffer + # to the host and validating its contents + ort_output_vals = ort_output.cpu().detach().numpy() + # Validate results + self.assertTrue(np.array_equal(torch_output_vals, ort_output_vals)) + def test_bind_input_and_non_preallocated_output(self): + torch_input, torch_output, torch_output_vals = self.create_model_and_input() + + session = onnxruntime.InferenceSession('model.onnx') + io_binding = session.io_binding() + + # Bind input to CUDA + io_binding.bind_input('input', torch_input.device.type, 0, np.float32, list(torch_input.size()), torch_input.data_ptr()) + + # Bind output to CUDA + # DISCLAIMER: This is only useful for ORT benchmarking as there is really no way to access the + # ORT allocated device memory for this bound output at this point as ORT doesn't provide + # a handle over this memory yet and to access the contents of this memory, we would have to copy contents + # over to the host (CPU) using copy_outputs_to_cpu() as done in the next steps. + # In future, we will provide a handle over the ORT allocated device memory for the bound output. + # For now, this is a useful tool for benchmarking as Run() doesn't include the device-host copy latency + # by doing this. + io_binding.bind_output('output', 'cuda') + + # Invoke Run + session.run_with_iobinding(io_binding) + + # Get outputs over to CPU (the outputs which were bound to CUDA will get copied over to the host here) + ort_output = io_binding.copy_outputs_to_cpu()[0] + + # Validate results + self.assertTrue(np.array_equal(torch_output_vals, ort_output)) + + + if __name__ == '__main__': unittest.main() diff --git a/orttraining/orttraining/python/ort_trainer.py b/orttraining/orttraining/python/ort_trainer.py index 33e00252e9..a928a1b149 100644 --- a/orttraining/orttraining/python/ort_trainer.py +++ b/orttraining/orttraining/python/ort_trainer.py @@ -102,7 +102,6 @@ def ort_training_session_run_helper(session, iobinding, inputs, input_descs, out torch_tensor = torch.zeros(output_desc.shape_, device=device, dtype=output_desc.eval_dtype_ if hasattr(output_desc, 'eval_dtype_') else output_desc.dtype_) - iobinding.bind_output(output_desc.name_, torch_tensor.device.type, get_device_index(device), dtype_torch_to_numpy(torch_tensor.dtype), list(torch_tensor.size()), torch_tensor.data_ptr()) @@ -195,6 +194,10 @@ def dtype_torch_to_numpy(torch_dtype): return np.int32 elif torch_dtype == torch.int16 or torch_dtype == torch.short: return np.int16 + elif torch_dtype == torch.bool: + return np.bool + else: + raise Exception("Torch type to numpy type mapping unavailable for: " + str(torch_dtype)) def wrap_for_input_match(model, loss_fn, input_names): import inspect diff --git a/tools/ci_build/build.py b/tools/ci_build/build.py index b0bf5388fd..ac34a5380c 100755 --- a/tools/ci_build/build.py +++ b/tools/ci_build/build.py @@ -1214,6 +1214,22 @@ def run_onnxruntime_tests(args, source_dir, ctest_path, build_dir, configs): [sys.executable, 'onnxruntime_test_python.py'], cwd=cwd, dll_path=dll_path) + iobinding_test = False + # For CUDA enabled builds test IOBinding feature + if args.use_cuda: + iobinding_test = True + try: + import torch # noqa + except ImportError as error: + iobinding_test = False + log.exception(error) + log.warning( + "Torch is not installed. " + "The IOBinding tests will be skipped as it requires Torch.") + + if iobinding_test: + run_subprocess([sys.executable, 'onnxruntime_test_python_iobinding.py'], cwd=cwd, dll_path=dll_path) + if not args.disable_ml_ops: run_subprocess([sys.executable, 'onnxruntime_test_python_mlops.py'], cwd=cwd, dll_path=dll_path) diff --git a/tools/ci_build/github/linux/docker/scripts/install_deps.sh b/tools/ci_build/github/linux/docker/scripts/install_deps.sh index 1f54119de1..eaa86c2fb7 100755 --- a/tools/ci_build/github/linux/docker/scripts/install_deps.sh +++ b/tools/ci_build/github/linux/docker/scripts/install_deps.sh @@ -113,8 +113,8 @@ if [ $DEVICE_TYPE = "Normal" ]; then ${PYTHON_EXE} -m pip install sympy==1.1.1 elif [ $DEVICE_TYPE = "gpu" ]; then ${PYTHON_EXE} -m pip install sympy==1.1.1 + ${PYTHON_EXE} -m pip install --upgrade --pre torch==1.6.0.dev20200610 torchvision==0.7.0.dev20200610 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html if [[ $BUILD_EXTR_PAR = *--enable_training* ]]; then - ${PYTHON_EXE} -m pip install --upgrade --pre torch==1.6.0.dev20200610 torchvision==0.7.0.dev20200610 -f https://download.pytorch.org/whl/nightly/cu101/torch_nightly.html ${PYTHON_EXE} -m pip install transformers==v2.10.0 # transformers requires sklearn ${PYTHON_EXE} -m pip install sklearn