hahnec / torchimize

numerical optimization in pytorch

Description

torchimize contains implementations of the Gradient Descent, Gauss-Newton and Levenberg-Marquardt optimization algorithms using the PyTorch library. The main motivation for this project is to enable convex optimization on GPUs based on the torch.Tensor class, which (as of 2022) is widely used in the deep learning field. This package features the capability to minimize several least-squares optimization problems at each loop iteration in parallel.

Installation

$ python3 -m pip install torchimize

Kick-Start

Single Cost Optimization

# single gradient descent
from torchimize.functions.single.gda_fun_single import gradient_descent
coeffs_list = gradient_descent(initials, cost_fun, args=(other_args,))

# single gauss-newton
from torchimize.functions import lsq_gna
coeffs_list = lsq_gna(initials, cost_fun, args=(other_args,))

# single levenberg-marquardt
from torchimize.functions import lsq_lma
coeffs_list = lsq_lma(initials, function=cost_fun, jac_function=jac_fun, args=(other_args,))

Parallel Cost Optimization

# parallel gradient descent for several optimization problems at multiple costs
from torchimize.functions import gradient_descent_parallel
coeffs_list = gradient_descent_parallel(
                    p = initials_batch,
                    function = multi_cost_fun_batch,
                    jac_function = multi_jac_fun_batch,
                    args = (other_args,),
                    wvec = torch.ones(5, device='cuda', dtype=initials_batch.dtype),
                    ftol = 1e-8,
                    ptol = 1e-8,
                    gtol = 1e-8,
                    l = 1.,
                    max_iter = 80,
                )

# parallel gauss-newton for several optimization problems at multiple costs
from torchimize.functions import lsq_gna_parallel
coeffs_list = lsq_gna_parallel(
                    p = initials_batch,
                    function = multi_cost_fun_batch,
                    jac_function = multi_jac_fun_batch,
                    args = (other_args,),
                    wvec = torch.ones(5, device='cuda', dtype=initials_batch.dtype),
                    ftol = 1e-8,
                    ptol = 1e-8,
                    gtol = 1e-8,
                    l = 1.,
                    max_iter = 80,
                )

# parallel levenberg-marquardt for several optimization problems at multiple costs
from torchimize.functions import lsq_lma_parallel
coeffs_list = lsq_lma_parallel(
                    p = initials_batch,
                    function = multi_cost_fun_batch,
                    jac_function = multi_jac_fun_batch,
                    args = (other_args,),
                    wvec = torch.ones(5, device='cuda', dtype=initials_batch.dtype),
                    ftol = 1e-8,
                    ptol = 1e-8,
                    gtol = 1e-8,
                    meth = 'marq',
                    max_iter = 40,
                )

# validate that your provided functions return correct tensor dimensionality
from torchimize.functions import test_fun_dims_parallel
ret = test_fun_dims_parallel(
    p = initials_batch,
    function = multi_cost_fun_batch,
    jac_function = multi_jac_fun_batch,
    args = (other_args,),
    wvec = torch.ones(5, device='cuda', dtype=initials_batch.dtype),
)

Note

For simultaneous minimization of B optimization problems at a multiple of C costs, the function and jac_function arguments require to return a torch.Tensor type of B x C x N and B x C x N x P, respectively. Here, N is the residual dimension and P represents the sought parameter number in each B x C.

For further details, see the API documentation.

Citation

@misc{torchimize,
    title={torchimize},
    author={Hahne, Christopher and Hayoz, Michel},
    year={2022},
    publisher = {GitHub},
    journal = {GitHub repository},
    howpublished = {\url{https://github.com/hahnec/torchimize}}
}