From 28188e1734c5b0f96ebb72408341fceb3d80aa06 Mon Sep 17 00:00:00 2001 From: Robert Parker Date: Fri, 1 Nov 2024 16:36:33 -0600 Subject: [PATCH 1/4] add gray_box_device option --- ext/MathOptAIPythonCallExt.jl | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl index 27559479..67280742 100644 --- a/ext/MathOptAIPythonCallExt.jl +++ b/ext/MathOptAIPythonCallExt.jl @@ -93,12 +93,15 @@ function MathOptAI.build_predictor( config::Dict = Dict{Any,Any}(), gray_box::Bool = false, gray_box_hessian::Bool = false, + gray_box_device = "cpu", ) if gray_box if !isempty(config) error("cannot specify the `config` kwarg if `gray_box = true`") end - return MathOptAI.GrayBox(predictor; hessian = gray_box_hessian) + return MathOptAI.GrayBox( + predictor; hessian = gray_box_hessian, device = gray_box_device + ) end torch = PythonCall.pyimport("torch") nn = PythonCall.pyimport("torch.nn") @@ -132,24 +135,25 @@ end function MathOptAI.GrayBox( predictor::MathOptAI.PytorchModel; hessian::Bool = false, + device::String = "cpu", ) torch = PythonCall.pyimport("torch") - torch_model = torch.load(predictor.filename) + torch_model = torch.load(predictor.filename).to(device) J = torch.func.jacrev(torch_model) H = torch.func.hessian(torch_model) # TODO(odow): I'm not sure if there is a better way to get the output # dimension of a torch model object? output_size(::Any) = PythonCall.pyconvert(Int, torch_model[-1].out_features) function callback(x) - py_x = torch.tensor(collect(x)) - py_value = torch_model(py_x).detach().numpy() + py_x = torch.tensor(collect(x), device=device) + py_value = torch_model(py_x).detach().cpu().numpy() value = PythonCall.pyconvert(Vector, py_value) - py_jacobian = J(py_x).detach().numpy() + py_jacobian = J(py_x).detach().cpu().numpy() jacobian = PythonCall.pyconvert(Matrix, py_jacobian) if !hessian return (; value, jacobian) end - hessians = PythonCall.pyconvert(Array, H(py_x).detach().numpy()) + hessians = PythonCall.pyconvert(Array, H(py_x).detach().cpu().numpy()) return (; value, jacobian, hessian = permutedims(hessians, (2, 3, 1))) end return MathOptAI.GrayBox(output_size, callback; has_hessian = hessian) From ca5b7d156cd4443197bf43da565a656f6a72fab2 Mon Sep 17 00:00:00 2001 From: Robert Parker Date: Tue, 5 Nov 2024 09:43:50 -0500 Subject: [PATCH 2/4] add gray_box_device option to docstring --- ext/MathOptAIPythonCallExt.jl | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl index 3b4cb546..6bb81a60 100644 --- a/ext/MathOptAIPythonCallExt.jl +++ b/ext/MathOptAIPythonCallExt.jl @@ -18,6 +18,8 @@ import MathOptAI config::Dict = Dict{Any,Any}(), reduced_space::Bool = false, gray_box::Bool = false, + gray_box_hessian::Bool = false, + gray_box_device::String = "cpu", ) Add a trained neural network from PyTorch via PythonCall.jl to `model`. @@ -41,6 +43,8 @@ Add a trained neural network from PyTorch via PythonCall.jl to `model`. nonlinear operator, with gradients provided by `torch.func.jacrev`. * `gray_box_hessian`: if `true`, the gray box additionally computes the Hessian of the output using `torch.func.hessian`. + * `gray_box_device`: device used to construct PyTorch tensors, e.g. `"cuda"` + to run on an Nvidia GPU. """ function MathOptAI.add_predictor( model::JuMP.AbstractModel, @@ -63,6 +67,7 @@ end config::Dict = Dict{Any,Any}(), gray_box::Bool = false, gray_box_hessian::Bool = false, + gray_box_device::String = "cpu", ) Convert a trained neural network from PyTorch via PythonCall.jl to a @@ -87,13 +92,15 @@ Convert a trained neural network from PyTorch via PythonCall.jl to a nonlinear operator, with gradients provided by `torch.func.jacrev`. * `gray_box_hessian`: if `true`, the gray box additionally computes the Hessian of the output using `torch.func.hessian`. + * `gray_box_device`: device used to construct PyTorch tensors, e.g. `"cuda"` + to run on an Nvidia GPU. """ function MathOptAI.build_predictor( predictor::MathOptAI.PytorchModel; config::Dict = Dict{Any,Any}(), gray_box::Bool = false, gray_box_hessian::Bool = false, - gray_box_device = "cpu", + gray_box_device::String = "cpu", ) if gray_box if !isempty(config) From 594911fd9455c9475dcc5508806095584d238d8f Mon Sep 17 00:00:00 2001 From: Robert Parker Date: Tue, 5 Nov 2024 09:48:26 -0500 Subject: [PATCH 3/4] format --- ext/MathOptAIPythonCallExt.jl | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl index 6bb81a60..92df3b6f 100644 --- a/ext/MathOptAIPythonCallExt.jl +++ b/ext/MathOptAIPythonCallExt.jl @@ -107,7 +107,9 @@ function MathOptAI.build_predictor( error("cannot specify the `config` kwarg if `gray_box = true`") end return MathOptAI.GrayBox( - predictor; hessian = gray_box_hessian, device = gray_box_device + predictor; + hessian = gray_box_hessian, + device = gray_box_device, ) end torch = PythonCall.pyimport("torch") @@ -145,14 +147,17 @@ function MathOptAI.GrayBox( device::String = "cpu", ) torch = PythonCall.pyimport("torch") - torch_model = torch.load(predictor.filename; weights_only = false).to(device) + torch_model = torch.load( + predictor.filename; + weights_only = false, + ).to(device) J = torch.func.jacrev(torch_model) H = torch.func.hessian(torch_model) # TODO(odow): I'm not sure if there is a better way to get the output # dimension of a torch model object? output_size(::Any) = PythonCall.pyconvert(Int, torch_model[-1].out_features) function callback(x) - py_x = torch.tensor(collect(x), device=device) + py_x = torch.tensor(collect(x), device = device) py_value = torch_model(py_x).detach().cpu().numpy() value = PythonCall.pyconvert(Vector, py_value) py_jacobian = J(py_x).detach().cpu().numpy() From 1d9e8f2e5c3e8881b453404771c8fb8a6459ad7e Mon Sep 17 00:00:00 2001 From: Robert Parker Date: Tue, 5 Nov 2024 09:57:06 -0500 Subject: [PATCH 4/4] format --- ext/MathOptAIPythonCallExt.jl | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl index 92df3b6f..cf0f1041 100644 --- a/ext/MathOptAIPythonCallExt.jl +++ b/ext/MathOptAIPythonCallExt.jl @@ -147,17 +147,15 @@ function MathOptAI.GrayBox( device::String = "cpu", ) torch = PythonCall.pyimport("torch") - torch_model = torch.load( - predictor.filename; - weights_only = false, - ).to(device) + torch_model = torch.load(predictor.filename; weights_only = false) + torch_model = torch_model.to(device) J = torch.func.jacrev(torch_model) H = torch.func.hessian(torch_model) # TODO(odow): I'm not sure if there is a better way to get the output # dimension of a torch model object? output_size(::Any) = PythonCall.pyconvert(Int, torch_model[-1].out_features) function callback(x) - py_x = torch.tensor(collect(x), device = device) + py_x = torch.tensor(collect(x); device = device) py_value = torch_model(py_x).detach().cpu().numpy() value = PythonCall.pyconvert(Vector, py_value) py_jacobian = J(py_x).detach().cpu().numpy()