Can gamspy solve neural network functions trained with pytorch as the objective function？ about gamspy HOT 9 CLOSED

from typing import Tuple

import gamspy as gp
import numpy as np


class ParameterMatrix:
    def __init__(
        self,
        model: gp.Container,
        name: str,
        data: np.ndarray,
        generate: bool = True,
        **kwargs
    ):
        self.model = model
        self.name = name
        self.shape = data.shape
        self.kwargs = kwargs
        self._domain = None
        self._symbol = None

        if generate:
            self._domain = []
            for size in self.shape:
                domain_name = f"md{size}"
                domain_symbol = None
                if domain_name not in self.model:
                    domain_symbol = gp.Set(model, name=domain_name, records=range(size))
                else:
                    domain_symbol = self.model[domain_name]

                self._domain.append(domain_symbol)

            self._symbol = gp.Parameter(model, name=name, domain=self._domain, records=data, **kwargs)

    def __matmul__(self, other):
        assert self.shape[1] == other.shape[0], "Matrix dims must match"
        ldomain = self._domain[0]
        mdomain = self._domain[1]
        rdomain = other._domain[1]
        return gp.Sum(
            mdomain,
            self._symbol[ldomain, mdomain] * other._symbol[mdomain, rdomain]
        )


class VariableMatrix:
    def __init__(
        self,
        model: gp.Container,
        name: str,
        shape: Tuple[int, int],
        generate: bool = True,
        **kwargs
    ):
        self.model = model
        self.name = name
        self.shape = shape
        self.kwargs = kwargs
        self._domain = None
        self._symbol = None

        if generate:
            self._domain = []
            for size in self.shape:
                domain_name = f"md{size}"
                domain_symbol = None
                if domain_name not in self.model:
                    domain_symbol = gp.Set(model, name=domain_name, records=range(size))
                else:
                    domain_symbol = self.model[domain_name]

                self._domain.append(domain_symbol)

            self._symbol = gp.Variable(model, name=name, domain=self._domain, **kwargs)

    def __matmul__(self, other):
        assert self.shape[1] == other.shape[0], "Matrix dims must match"
        ldomain = self._domain[0]
        mdomain = self._domain[1]
        rdomain = other._domain[1]
        return gp.Sum(
            mdomain,
            self._symbol[ldomain, mdomain] * other._symbol[mdomain, rdomain]
        )

import torch
import torch.nn as nn

import uuid
import gamspy as gp
import gamspy.math
import numpy as np

from matrix import ParameterMatrix, VariableMatrix

def get_auto_name(prefix):
    name = str(uuid.uuid4()).split("-")[0]
    return f"{prefix}_{name}"


def ReLU(m: gp.Container, expr, expr_domain, expr_shape):
    name = get_auto_name("ReLU")
    var = gp.Variable(m, name=name, domain=expr_domain, type="positive")

    eq1_name = get_auto_name("calc_ReLU_1")
    eq1 = gp.Equation(m, name=eq1_name, domain=expr_domain)
    eq1[...] = var[...] >= expr[...]

    binary_var_name = get_auto_name("b_ReLU")
    binary_var = gp.Variable(m, name=binary_var_name, domain=expr_domain, type="binary")

    eq2_name = get_auto_name("calc_ReLU_2")
    eq2 = gp.Equation(m, name=eq2_name, domain=expr_domain)
    eq2[...] = var[...] <= expr[...] * binary_var[...]

    var_matrix = VariableMatrix(m, name=name, shape=expr_shape, generate=False)
    var_matrix._symbol = var
    var_matrix._domain = var.domain

    return var_matrix


class NNModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(NNModel, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        return out

input_size = 2
hidden_size = 64
output_size = 1

model = NNModel(input_size, hidden_size, output_size)
model.load_state_dict(torch.load("model.pth"))

m = gp.Container()

x1 = VariableMatrix(m, "x_input", shape=(input_size, 1))

# Create a parameter matrix for w1.T
w1 = ParameterMatrix(m, "w1t", model.fc1.weight.detach().numpy())

# Create par for bias 1
b1 = ParameterMatrix(m, "b1", model.fc1.bias.detach().numpy())

# Create a parameter matrix for w2.T
w2 = ParameterMatrix(m, "w2t", model.fc2.weight.detach().numpy())

# Create par for bias 2
b2 = ParameterMatrix(m, "b2", model.fc2.bias.detach().numpy())

out1 = VariableMatrix(m, "out1", shape=(hidden_size, 1)) # after fc1
out3 = VariableMatrix(m, "out3", shape=(output_size, 1)) # after fc2

eq_calc_out_1 = gp.Equation(m, name="calc_out_1", domain=out1._domain)
eq_calc_out_1[...] = out1._symbol[...] == w1 @ x1 + b1._symbol[...]


out2 = ReLU(m, out1._symbol, out1._domain, (hidden_size, 1))

eq_calc_out_3 = gp.Equation(m, name="calc_out_3", domain=out3._domain)
eq_calc_out_3[...] = out3._symbol[...] == w2 @ out2 + b2._symbol[...]


# Let's set x1 to something to see if it works:
x1._symbol.fx[...] = 2

z = gp.Variable(m, name="fake_obj")

eq_obj = gp.Equation(m, name="calc_obj")
eq_obj[...] = z == out3._symbol["0", "0"]


model_fit = gp.Model(
    m,
    name="fit",
    equations=m.getEquations(),
    problem="MINLP",
    sense="min",
    objective=z,
)
model_fit.solve()


val = torch.tensor([2.0, 2.0])
print(model(val))
print(z.records)