This is the code of paper "A Deep Instance Generative Framework for MILP Solvers Under Limited Data Availability". Zijie Geng, Xijun Li, Jie Wang, Xiao Li, Yongdong Zhang, Feng Wu. NeurIPS 2023 (Spotlight).
import os.path as path

import hydra
import torch
from omegaconf import DictConfig

from src import G2MILP, Benchmark, Generator, set_cpu_num, set_seed


@hydra.main(version_base=None, config_path="conf", config_name="generate")
def generate(config: DictConfig):
    """
    Generate instances using G2MILP.
    """
    set_seed(config.seed)
    set_cpu_num(config.num_workers + 1)
    torch.cuda.set_device(config.cuda)

    model_path = path.join(config.paths.model_dir, "model_best.ckpt")
    model = G2MILP.load_model(config, model_path)

    generator = Generator(
        model=model,
        config=config.generator,
        templates_dir=config.paths.dataset_samples_dir,
        save_dir=config.paths.samples_dir,
    )
    generator.generate()

    benchmarker = Benchmark(
        config=config.benchmarking,
        dataset_stats_dir=config.paths.dataset_stats_dir,
    )
    results = benchmarker.assess_samples(
        samples_dir=config.paths.samples_dir,
        benchmark_dir=config.paths.benchmark_dir,
    )

    info_path = path.join(config.paths.benchmark_dir, "info.json")
    benchmarker.log_info(
        generator_config=config.generator,
        benchmarking_config=config.benchmarking,
        meta_results=results,
        save_path=info_path,
    )


if __name__ == '__main__':
    generate()

generate.py
MIRALab-USTC/L2O-G2MILP
import json
import logging
import multiprocessing as mp
import os
import os.path as path
import pickle
from functools import partial
from typing import List

import hydra
import numpy as np
from omegaconf import DictConfig
from pandas import DataFrame
from tqdm import tqdm

from src import instance2graph, set_cpu_num, set_seed, solve_instance


def preprocess_(file: str, config: DictConfig):
    """
    Preprocesses a single instance.

    Args:
        file: instance file name
        config: config
    """
    sample_path = path.join(config.paths.data_dir, file)
    data, features = instance2graph(sample_path, config.compute_features)

    with open(path.join(config.paths.dataset_samples_dir, os.path.splitext(file)[0]+".pkl"), "wb") as f:
        pickle.dump(data, f)
    if config.solve_instances:
        solving_results = {"instance": file}
        solving_results.update(solve_instance(sample_path))
    else:
        solving_results = None
    return features, solving_results

def make_dataset_features(
        features: List[dict],
        solving_results: List[dict],
        config:DictConfig
    ):
    """
    Computes the dataset features.

    Args:
        features: list of instance features
        solving_results: list of solving results
        config: config
    """
    if config.compute_features:
        logging.info(f"Writing instance features to: {config.paths.dataset_features_path}")
        features: DataFrame = DataFrame(features, columns=features[0].keys()).set_index("instance")
        features.to_csv(config.paths.dataset_features_path)

        logging.info(f"Writing dataset statistics to: {config.paths.dataset_stats_path}")
        stats = {
            "rhs_type": config.dataset.rhs_type,
            "rhs_min": np.min(features["rhs_min"]),
            "rhs_max": np.max(features["rhs_max"]),

            "obj_type": config.dataset.obj_type,
            "obj_min": np.min(features["obj_min"]),
            "obj_max": np.max(features["obj_max"]),
            
            "coef_type": config.dataset.lhs_type,
            "coef_min": np.min(features['lhs_min']),
            "coef_max": np.max(features["lhs_max"]),
            "coef_dens": np.mean(features["coef_dens"]),

            "cons_degree_min": int(np.min(features["cons_degree_min"])),
            "cons_degree_max": int(np.max(features["cons_degree_max"])),
        }
        with open(config.paths.dataset_stats_path, "w") as f:
            f.write(json.dumps(stats, indent=2))

    if config.solve_instances:
        logging.info(f"Writting solving results to: {config.paths.dataset_solving_path}")
        solving_results: DataFrame = DataFrame(solving_results, columns=solving_results[0].keys()).set_index("instance")
        solving_results.to_csv(config.paths.dataset_solving_path)

        solving_time = solving_results.loc[:, ["solving_time"]].to_numpy()
        num_nodes = solving_results.loc[:, ["num_nodes"]].to_numpy()
        
        logging.info(f"  mean solving time: {solving_time.mean()}")
        logging.info(f"  mean num nodes: {num_nodes.mean()}")

    
@hydra.main(version_base=None, config_path="conf", config_name="preprocess")
def preprocess(config: DictConfig):
    """
    Preprocesses the dataset.
    """
    set_seed(config.seed)
    set_cpu_num(config.num_workers + 1)

    logging.info("="*10 + "Begin preprocessing" + "="*10)
    logging.info(f"Dataset: {config.dataset.name}.")
    logging.info(f"Dataset dir: {config.paths.data_dir}")

    os.makedirs(config.paths.dataset_samples_dir, exist_ok=True)
    os.makedirs(config.paths.dataset_stats_dir, exist_ok=True)

    files: list = os.listdir(config.paths.data_dir)
    files.sort()
    if len(files) > config.dataset.num_train:
        files = files[:config.dataset.num_train]

    func = partial(preprocess_, config=config)
    with mp.Pool(config.num_workers) as pool:
        features, solving_results = zip(*list(tqdm(pool.imap(func, files), total=len(files), desc="Preprocessing")))
    logging.info(f"Preprocessed samples are saved to: {config.paths.dataset_samples_dir}")

    make_dataset_features(features, solving_results, config)
    
    logging.info("="*10 + "Preprocessing finished" + "="*10)


if __name__ == "__main__":
    preprocess()

preprocess.py
MIRALab-USTC/L2O-G2MILP
from src.benchmarks import Benchmark
from src.data import InstanceDataset, TrainSampler
from src.generator import Generator
from src.model import G2MILP
from src.trainer import Trainer
from src.utils import instance2graph, set_cpu_num, set_seed, solve_instance

src/__init__.py
MIRALab-USTC/L2O-G2MILP
from src.benchmarks.benchmarks import Benchmark
src/benchmarks/__init__.py
MIRALab-USTC/L2O-G2MILP
import json
import logging
import os

import pandas as pd
from omegaconf import DictConfig, OmegaConf


from .utils import FEATURES, compute_features, compute_jsdiv, solve_instances


class Benchmark():
    def __init__(self, config: DictConfig, dataset_stats_dir: str):
        self.config = config
        self.dataset_stats_dir = dataset_stats_dir

    def assess_samples(self,
                       samples_dir: str,
                       benchmark_dir: str
                       ):
        """
        Assess the generated samples.

        Args:
            samples_dir: directory of generated samples
            benchmark_dir: directory to save benchmark results

        Returns:
            results: dict, benchmark results
        """
        os.makedirs(benchmark_dir, exist_ok=True)
        distribution_results = assess_distribution(
            self.config, samples_dir, self.dataset_stats_dir, benchmark_dir)
        solving_results = assess_solving_results(
            self.config, samples_dir, self.dataset_stats_dir, benchmark_dir)
        results = {
            "distribution": distribution_results,
            "solving": solving_results,
        }
        return results

    @staticmethod
    def log_info(
        generator_config: DictConfig,
        benchmarking_config: DictConfig,
        meta_results: dict,
        save_path: str
    ):
        info = {
            "generator": OmegaConf.to_container(generator_config),
            "benchmarking": OmegaConf.to_container(benchmarking_config),
            "meta_results": meta_results,
        }
        logging.info("-"*10 + " Benchmarking  Info " + "-"*10)
        json_msg = json.dumps(info, indent=4)
        for msg in json_msg.split("\n"):
            logging.info(msg)
        logging.info("-"*40)

        if save_path is not None:
            with open(save_path, 'w') as f:
                f.write(json.dumps(info, indent=4))
            logging.info(f"Benchmarking info saved to {save_path} .")


def assess_distribution(config, samples_dir, dataset_stats_dir, benchmark_dir):
    samples_features = compute_features(
        samples_dir, num_workers=config.num_workers)
    samples_features = pd.DataFrame(samples_features).set_index("instance")
    samples_features.to_csv(os.path.join(benchmark_dir, "features.csv"))

    reference_features = pd.read_csv(os.path.join(
        dataset_stats_dir, "features.csv")).set_index("instance")

    used_features = list(FEATURES.keys())
    samples_features = samples_features.loc[:, used_features].to_numpy()
    reference_features = reference_features.loc[:, used_features].to_numpy()

    score, meta_results = compute_jsdiv(
        samples_features, reference_features, num_samples=config.num_samples)

    results = {
        "score": score,
        "meta_results": meta_results,
    }
    return results


def assess_solving_results(config, samples_dir, dataset_stats_dir, benchmark_dir):
    samples_solving_results = solve_instances(
        samples_dir, num_workers=config.num_workers)
    samples_solving_results = pd.DataFrame(
        samples_solving_results).set_index("instance")
    samples_solving_results.to_csv(os.path.join(
        benchmark_dir, "solving_results.csv"))

    reference_solving_results = pd.read_csv(os.path.join(
        dataset_stats_dir, "solving_results.csv")).set_index("instance")

    samples_solving_time = samples_solving_results.loc[:, [
        "solving_time"]].to_numpy()
    samples_num_nodes = samples_solving_results.loc[:, [
        "num_nodes"]].to_numpy()
    reference_solving_time = reference_solving_results.loc[:, [
        "solving_time"]].to_numpy()
    reference_num_nodes = reference_solving_results.loc[:, [
        "num_nodes"]].to_numpy()

    results = {
        "solving_time": {
            "mean": samples_solving_time.mean(),
            "mean_reference": reference_solving_time.mean(),
            "mean_error": abs(samples_solving_time.mean() - reference_solving_time.mean()) / reference_solving_time.mean(),
        },
        "num_nodes": {
            "mean": samples_num_nodes.mean(),
            "mean_reference": reference_num_nodes.mean(),
            "mean_error": abs(samples_num_nodes.mean() - reference_num_nodes.mean()) / reference_num_nodes.mean(),
        },
    }
    return results

src/benchmarks/benchmarks.py
MIRALab-USTC/L2O-G2MILP
import multiprocessing as mp
import os
import os.path as path
from functools import partial
from typing import Union

import numpy as np
from scipy.stats import entropy
from tqdm import tqdm

from src.utils import instance2graph, solve_instance

FEATURES = {
    "coef_dens": float,
    "var_degree_mean": float,
    "var_degree_std": float,
    "cons_degree_mean": float,
    "cons_degree_std": float,
    "lhs_mean": float,
    "lhs_std": float,
    "rhs_mean": float,
    "rhs_std": float,

    "clustering": float,
    "modularity": float,
}


def js_div(x1: np.ndarray, x2: np.ndarray) -> float:
    x = np.hstack([x1, x2])
    if x.std() < 1e-10:
        return 0.0
    M, bins = np.histogram(x, bins=5, density=True)

    P, _ = np.histogram(x1, bins=bins)
    Q, _ = np.histogram(x2, bins=bins)

    return (entropy(P, M) + entropy(Q, M)) / 2


def compute_jsdiv(features1: Union[np.ndarray, str], features2: Union[np.ndarray, str], num_samples: int = 1000):

    f1 = features1[np.random.choice(
        list(range(len(features1))), num_samples, replace=True), :]
    f2 = features2[np.random.choice(
        list(range(len(features2))), num_samples, replace=True), :]

    meta_results = {}
    for i in range(len(FEATURES)):
        jsdiv = js_div(f1[:, i], f2[:, i])
        meta_results[list(FEATURES.keys())[i]] = round(
            1 - jsdiv / np.log(2), 3)

    score = sum(meta_results.values()) / len(meta_results)
    return score, meta_results


def compute_features(samples_dir: str, num_workers: int = 1):
    samples = os.listdir(samples_dir)

    func = partial(compute_features_, data_dir=samples_dir)
    with mp.Pool(num_workers) as pool:
        features = list(tqdm(pool.imap(func, samples), total=len(
            samples), desc="Computing features"))

    return features


def compute_features_(file: str, data_dir: str):
    sample_path = path.join(data_dir, file)
    _, features = instance2graph(sample_path, compute_features=True)
    return features


def solve_instances(samples_dir: str, num_workers: int):
    samples = os.listdir(samples_dir)

    func = partial(solve_instance_, samples_dir=samples_dir)
    with mp.Pool(num_workers) as pool:
        solving_results = list(
            tqdm(pool.imap(func, samples), total=len(samples), desc="Solving instances"))

    return solving_results


def solve_instance_(instance: str, samples_dir: str):
    sample_path = os.path.join(samples_dir, instance)
    solving_results = {"instance": instance}
    solving_results.update(solve_instance(sample_path))
    return solving_results

src/benchmarks/utils.py
MIRALab-USTC/L2O-G2MILP
import os
import pickle

import pandas as pd
import torch
from torch import FloatTensor, LongTensor
from torch.utils.data import BatchSampler
from torch_geometric.data import Data, Dataset


class BipartiteGraph(Data):
    masked_cons_idx: int
    bias_label: FloatTensor
    degree_label: LongTensor
    logits_label: FloatTensor
    connected_vars_idx: LongTensor
    weights_label: FloatTensor

    def __init__(
        self,
        x_constraints: FloatTensor = None,
        edge_index: LongTensor = None,
        edge_attr: FloatTensor = None,
        x_variables: FloatTensor = None,
        **kwargs
    ):
        """
        Bipartite graph data object, each representing a MILP instance.
        """
        super().__init__(**kwargs)
        self.x_constraints = x_constraints
        self.edge_index = edge_index
        self.edge_attr = edge_attr
        self.x_variables = x_variables
        if x_constraints is not None:
            self.num_nodes = len(x_constraints) + len(x_variables)

    def __inc__(self, key, value, store, *args, **kwargs):
        if key == "edge_index":
            return torch.tensor(
                [[self.x_constraints.size(0)], [self.x_variables.size(0)]]
            )
        if key == "masked_cons_idx":
            return self.x_constraints.size(0)
        elif key == "connected_vars_idx":
            return self.x_variables.size(0)
        else:
            return super().__inc__(key, value, *args, **kwargs)

    @property
    def num_variables(self):
        return len(self.x_variables)

    @property
    def num_constraints(self):
        return len(self.x_constraints)


class InstanceDataset(Dataset):
    def __init__(self,
                 data_dir: str,
                 solving_results_path: str = None,
                 num_instances: int = None,
                 ):
        """
        Dataset for instances.

        Args:
            data_dir: path to the directory saving the instances
            solving_results_path: path to the csv file saving the solving results
            num_instances: number of instances to use
        """
        super().__init__()
        self.data_dir = data_dir
        self.files = sorted(os.listdir(data_dir))
        self.solving_path = solving_results_path

        if self.solving_path:
            self.solving_results = pd.read_csv(
                solving_results_path).set_index("instance")
            self.solving_time_mean = self.solving_results.loc[:, "solving_time"].mean(
            )
            self.solving_time_std = self.solving_results.loc[:, "solving_time"].std(
            )
            self.num_nodes_mean = self.solving_results.loc[:, "num_nodes"].mean(
            )
            self.num_nodes_std = self.solving_results.loc[:, "num_nodes"].std()

        if num_instances:
            self.files = self.files[:num_instances]

    def len(self):
        return len(self.files)

    def get(self, index):
        file = os.path.join(self.data_dir, self.files[index])
        with open(file, "rb") as f:
            data = pickle.load(f)
        x_constraints, edge_index, edge_attr, x_variables = data

        return BipartiteGraph(
            x_constraints=FloatTensor(x_constraints),
            edge_index=LongTensor(edge_index),
            edge_attr=FloatTensor(edge_attr),
            x_variables=FloatTensor(x_variables),
        )


class TrainSampler(BatchSampler):
    def __init__(self, dataset: InstanceDataset, batch_size: int, repeat_size: int, total_steps: int):
        """
        Sampler for training G2MILP.

        Args:
            dataset: training dataset
            batch_size: batch size, number of instances in each batch
            repeat_size: number of times to repeat randomly masking constraints of the same instance. We use this to increase the training difficulty
            total_steps: number of total training steps
        """
        self.dataset = dataset
        self.batch_size = batch_size
        self.repeat_size = repeat_size
        self.total_steps = total_steps

    def __iter__(self):
        all_batches = torch.randint(
            0, len(self.dataset), (self.total_steps, self.batch_size))
        for i in range(self.total_steps):
            batch = all_batches[i]
            batch = batch.repeat_interleave(self.repeat_size)
            batch = batch.tolist()
            yield batch

    def __len__(self):
        return self.total_steps

src/data.py
MIRALab-USTC/L2O-G2MILP
import logging
import os
import os.path as path
from typing import Iterator, List

import torch
from omegaconf import DictConfig
from torch.utils.data import BatchSampler, DataLoader, RandomSampler

from src.data import BipartiteGraph, InstanceDataset
from src.model import G2MILP
from src.utils import graph2instance


class Generator():
    def __init__(self,
                 model: G2MILP,
                 config: DictConfig,
                 templates_dir: str,
                 save_dir: str,
                 ):
        """
        Generator wrapper for G2MILP.

        Args:
            model: G2MILP model
            config: generate config
            templates_dir: directory of template instances
            save_dir: directory to save generated instances
        """
        self.model = model
        self.templates_dir = templates_dir
        self.config = config
        self.samples_dir = save_dir

    def generate(self):
        os.makedirs(self.samples_dir, exist_ok=True)
        template_dataset = InstanceDataset(self.templates_dir)
        template_loader: Iterator[List[BipartiteGraph]] = DataLoader(
            dataset=template_dataset,
            batch_sampler=BatchSampler(
                sampler=RandomSampler(
                    template_dataset, replacement=True, num_samples=self.config.num_samples),
                batch_size=self.config.batch_size,
                drop_last=False),
            collate_fn=lambda x: [x_.cuda() for x_ in x]
        )

        i = 0
        self.model.eval()
        self.model.zero_grad()

        logging.info(
            "="*10 + f"Generating {self.config.num_samples} instances" + "="*10)
        with torch.no_grad():
            for batch in template_loader:
                sample_graphs = self.model.decode(batch, self.config)
                for sample_graph in sample_graphs:
                    i += 1
                    sample_model = graph2instance(sample_graph)
                    sample_path = path.join(
                        self.samples_dir, f"instance_{i}.lp")
                    sample_model.writeProblem(sample_path)
        logging.info("="*40)

src/generator.py
MIRALab-USTC/L2O-G2MILP
import json
from typing import List, Union

import torch
import torch.nn as nn
from omegaconf import DictConfig
from torch import Tensor
from torch_geometric.data import Batch
from torch_geometric.utils import sort_edge_index

import src.tb_writter as tb_writter
from src.data import BipartiteGraph
from src.modules import (Bias_Predictor, Degree_Predictor, Logits_Predictor,
                         ReSample, Weights_Predictor)
from src.nn import BipartiteGraphEmbedding, BipartiteGraphGNN


class G2MILP(nn.Module):

    def __init__(self,
                 config: DictConfig,
                 data_stats: dict
                 ):
        """
        G2MILP model, the main model.

        Args:
            config: model config
            data_stats: dataset statistics, used to configure the model
        """
        super().__init__()

        self.config = config

        self.embedding_layer = BipartiteGraphEmbedding(config.graph_embedding)
        self.encoder_layer = BipartiteGraphGNN(config.gnn)
        self.decoder_layer = BipartiteGraphGNN(config.gnn, is_masked=True)
        self.resample_layer = ReSample(config.resample)

        self.bias_predictor = Bias_Predictor(
            config.bias_predictor,
            data_stats,
        )
        self.degree_predictor = Degree_Predictor(
            config.degree_predictor,
            data_stats,
        )
        self.logits_predictor = Logits_Predictor(
            config.logits_predictor,
            data_stats,
        )
        self.weights_predictor = Weights_Predictor(
            config.weights_predictor,
            data_stats,
        )

    def forward(self,
                data: Union[Batch, BipartiteGraph],
                beta_cons: float,
                beta_var: float
                ) -> Tensor:
        """
        Forward pass of the model (for training).

        Inputs:
            data: batch of bipartite graphs
            beta_cons: coefficient of the KL loss for constraints
            beta_var: coefficient of the KL loss for variables

        Outputs:
            loss: loss of the model
        """
        orig_graph = self.embedding_layer.embed_graph(data)
        masked_graph = self.embedding_layer.mask_and_embed_graph(data)

        z_conss, z_vars = self.encoder_layer.forward(orig_graph)
        z_conss, z_vars, cons_kl_loss, var_kl_loss = self.resample_layer.forward(
            z_conss, z_vars)

        h_conss, h_vars = self.decoder_layer.forward(masked_graph)

        cons_loss, cons_pred = self.bias_predictor.forward(
            z_conss, h_conss, masked_graph.masked_cons_idx, masked_graph.bias_label)

        degree_loss, degree_pred = self.degree_predictor.forward(
            z_conss, h_conss, masked_graph.masked_cons_idx, masked_graph.degree_label)

        logits_loss, logits_pred = self.logits_predictor.forward(
            z_vars, h_vars, masked_graph.logits_label)

        weights_loss, weights_pred = self.weights_predictor.forward(
            z_vars, h_vars, masked_graph.connected_vars_idx, masked_graph.weights_label)

        cons_loss = cons_loss * self.config.loss_weights.cons_loss
        degree_loss = degree_loss * self.config.loss_weights.degree_loss
        logits_loss = logits_loss * self.config.loss_weights.logits_loss
        weights_loss = weights_loss * self.config.loss_weights.weights_loss
        loss = (beta_cons * cons_kl_loss + beta_var * var_kl_loss + cons_loss +
                degree_loss + logits_loss + weights_loss) / data.num_graphs

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_histogram("Embeddings/h_vars", h_vars, tb_writter.step)
            tb.add_histogram("Embeddings/h_conss", h_conss, tb_writter.step)
            tb.add_histogram("Embeddings/z_vars", z_vars, tb_writter.step)
            tb.add_histogram("Embeddings/z_conss", z_conss, tb_writter.step)
        else:
            tb.add_histogram("Embeddings/h_vars_val", h_vars, tb_writter.step)
            tb.add_histogram("Embeddings/h_conss_val",
                             h_conss, tb_writter.step)
            tb.add_histogram("Embeddings/z_vars_val", z_vars, tb_writter.step)
            tb.add_histogram("Embeddings/z_conss_val",
                             z_conss, tb_writter.step)

        return loss

    def decode(self, graphs: List[BipartiteGraph], config: dict):

        avg_num_constraints = sum([len(graph.x_constraints)
                                  for graph in graphs]) / len(graphs)
        num_iters = round(avg_num_constraints * config.mask_ratio)
        for _ in range(num_iters):

            masked_graphs = self.embedding_layer.mask_and_embed_graph(graphs)

            z_conss = torch.randn(
                (len(masked_graphs.x_constraints_batch), self.config.common.embd_size)).cuda()
            z_vars = torch.randn(
                (len(masked_graphs.x_variables_batch), self.config.common.embd_size)).cuda()

            h_conss, h_vars = self.decoder_layer(masked_graphs)

            conss = self.bias_predictor.decode(
                z_conss, h_conss, masked_graphs.masked_cons_idx)

            degree = self.degree_predictor.decode(
                z_conss, h_conss, masked_graphs.masked_cons_idx)

            weights_idx = self.logits_predictor.decode(
                z_vars, h_vars, masked_graphs.x_variables_batch, degree)

            weights = self.weights_predictor.decode(
                z_vars, h_vars, masked_graphs.x_variables_batch, weights_idx)

            masked_graphs = masked_graphs.to_data_list()
            weights_ptr = 0
            for bid, (graph, masked_graph) in enumerate(zip(graphs, masked_graphs)):
                masked_cons_idx = masked_graph.masked_cons_idx.cuda()
                graph.x_constraints[masked_cons_idx] = conss[bid].view(-1, 1)

                seen_edges = torch.where(
                    graph.edge_index[0] != masked_cons_idx)[0]
                new_edge_index = torch.LongTensor(
                    [[masked_cons_idx, i] for i in weights_idx[bid]]).T.cuda()
                new_edge_attr = weights[weights_ptr: weights_ptr +
                                        new_edge_index.shape[1]]
                edge_index = torch.cat(
                    (graph.edge_index[:, seen_edges], new_edge_index), dim=-1)
                edge_attr = torch.cat(
                    (graph.edge_attr[seen_edges], new_edge_attr), dim=0)
                graph.edge_index, graph.edge_attr = sort_edge_index(
                    edge_index, edge_attr)

                graphs[bid] = graph
                weights_ptr += len(new_edge_index)

        results = []
        for graph in graphs:
            x_constraints = graph.x_constraints.detach().cpu().numpy()
            edge_index = graph.edge_index.detach().cpu().numpy()
            edge_attr = graph.edge_attr.detach().cpu().numpy()
            x_variables = graph.x_variables.detach().cpu().numpy()
            results.append([x_constraints, edge_index, edge_attr, x_variables])
        return results

    @staticmethod
    def load_model(config: DictConfig, load_model_path: str = None) -> "G2MILP":
        """
        Loads the model.
        """
        data_stats = json.load(open(config.paths.dataset_stats_path, 'r'))
        model = G2MILP(config.model, data_stats).cuda()

        if load_model_path:
            load_ckpt = torch.load(load_model_path)
            model.load_state_dict(load_ckpt, strict=False)

        return model
src/model.py
MIRALab-USTC/L2O-G2MILP
import torch
import torch.nn as nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
import src.tb_writter as tb_writter
from sklearn.metrics import roc_auc_score
from torch_geometric.utils import unbatch
from src.utils import downsample
from torch import Tensor

class ReSample(nn.Module):
    def __init__(self, config) -> None:
        """
        The resample layer in VAE.
        """
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.z_conss_mean = nn.Sequential(
            nn.Linear(self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.embd_size),
        )
        self.z_vars_mean = nn.Sequential(
            nn.Linear(self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.embd_size),
        )
      
        self.z_conss_logstd = nn.Sequential(
            nn.Linear(self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.embd_size),
        )
        self.z_vars_logstd = nn.Sequential(
            nn.Linear(self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.embd_size),
        )
    
    def forward(self,
            z_conss: Tensor,
            z_vars: Tensor
        ):
        """
        Input:
            z_conss: (num_cons, embd_size), constraint embeddings
            z_vars: (num_vars, embd_size), variable embeddings

        Output:
            z_conss: (num_cons, embd_size), 
            z_vars: (num_vars, embd_size)
            conss_kl_loss: scalar
            vars_kl_loss: scalar
        """
        z_conss_mean = self.z_conss_mean(z_conss)
        z_conss_logstd = torch.clamp_max(self.z_conss_logstd(z_conss), max=10)
        conss_kl_loss = - 0.5 * torch.sum(1.0 + z_conss_logstd - z_conss_mean * z_conss_mean - torch.exp(z_conss_logstd))
        if self.training:
            conss_epsilon = torch.randn_like(z_conss_mean, device=z_conss.device)
            z_conss = z_conss_mean + torch.exp(z_conss_logstd / 2) * conss_epsilon
        else:
            z_conss = z_conss_mean

        z_vars_mean = self.z_vars_mean(z_vars)
        z_vars_logstd = torch.clamp_max(self.z_vars_logstd(z_vars), max=10)
        vars_kl_loss = - 0.5 * torch.sum(1.0 + z_vars_logstd - z_vars_mean * z_vars_mean - torch.exp(z_vars_logstd))
        if self.training:
            vars_epsilon = torch.randn_like(z_vars_mean, device=z_vars.device)
            z_vars = z_vars_mean + torch.exp(z_vars_logstd / 2) * vars_epsilon
        else:
            z_vars = z_vars_mean

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("resample/conss_kl_loss", conss_kl_loss, tb_writter.step)
            tb.add_scalar("resample/vars_kl_loss", vars_kl_loss, tb_writter.step)
        else:
            tb.add_scalar("resample/conss_kl_loss_val", conss_kl_loss, tb_writter.step)
            tb.add_scalar("resample/vars_kl_loss_val", vars_kl_loss, tb_writter.step)

        return z_conss, z_vars, conss_kl_loss, vars_kl_loss

class Obj_Predictor(nn.Module):
    def __init__(self, config, data_stats):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.obj_type = data_stats["obj_type"]
        self.hidden_size = self.hidden_size

        self.obj_min = int(data_stats["obj_min"])
        self.obj_max = int(data_stats["obj_max"])

        self.var_obj_predictor = nn.Sequential(
            nn.Linear(2 * self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1),
            nn.Sigmoid(),
        )
       
        self.var_obj_loss = MSELoss(reduction="sum")
    
    def convert(self, obj: torch.Tensor):
        if self.obj_type == "int":
            return torch.round(obj * (self.obj_max - self.obj_min) + self.obj_min)
        elif self.obj_type == "float":
            return obj * (self.obj_max - self.obj_min) + self.obj_min

    def forward(self, z_vars: torch.Tensor, h_vars: torch.Tensor, masked_var_idx: torch.LongTensor, obj_label: torch.Tensor):
        inputs = torch.cat((z_vars[masked_var_idx], h_vars[masked_var_idx]), dim=-1)
        obj_pred = self.var_obj_predictor(inputs).view(-1)

        if abs(self.obj_max - self.obj_min) < 1e-6:
            obj_loss = torch.tensor(0.0, device=inputs.device)
        else:
            obj_label_ = (obj_label.view(-1) - self.obj_min) / (self.obj_max - self.obj_min)
            obj_loss = self.var_obj_loss(obj_pred, obj_label_)
    
        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("Variable_predictor/obj_loss", obj_loss, tb_writter.step)
            tb.add_histogram("Variable_predictor/obj", obj_label, tb_writter.step)
            tb.add_histogram("Variable_predictor/obj_pred", self.convert(obj_pred), tb_writter.step)
        else:
            tb.add_scalar("Variable_predictor/obj_loss_val", obj_loss, tb_writter.step)
            tb.add_histogram("Variable_predictor/obj_val", obj_label, tb_writter.step)
            tb.add_histogram("Variable_predictor/obj_pred_val", self.convert(obj_pred), tb_writter.step)

        return obj_loss, obj_pred

    def decode(self, z_vars: torch.Tensor, h_vars: torch.Tensor, masked_var_idx: torch.LongTensor):
        inputs = torch.cat((z_vars[masked_var_idx], h_vars[masked_var_idx]), dim=-1)
        obj_pred = self.var_obj_predictor(inputs).view(-1)
        return self.convert(obj_pred)

class Type_Predictor(nn.Module):
    def __init__(self, config, data_stats):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.var_type_predictor = nn.Sequential(
            nn.Linear(2 * self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 4)
        )

        self.var_type_loss = CrossEntropyLoss(reduction="sum")
    
    def convert(self, type: torch.Tensor):
        type = torch.argmax(type, dim=1)
        return torch.zeros((len(type), 4), device=type.device).scatter_(1, type.view(-1,1), 1)

    def forward(self, z_vars: torch.Tensor, h_vars: torch.Tensor, masked_var_idx: torch.LongTensor, type_label: torch.LongTensor):
        inputs = torch.cat((z_vars[masked_var_idx], h_vars[masked_var_idx]), dim=-1)
    
        type_pred = self.var_type_predictor(inputs)
        type_loss = self.var_type_loss(type_pred, type_label.view(-1))

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("Variable_predictor/type_loss", type_loss, tb_writter.step)
            tb.add_histogram("Variable_predictor/type", type_label, tb_writter.step)
            tb.add_histogram("Variable_predictor/type_pred", self.convert(type_pred), tb_writter.step)
        else:
            tb.add_histogram("Variable_predictor/type_val", type_label, tb_writter.step)
            tb.add_histogram("Variable_predictor/type_pred_val", self.convert(type_pred), tb_writter.step)

        return type_loss, type_pred

    def decode(self, z_vars: torch.Tensor, h_vars: torch.Tensor, masked_var_idx: torch.LongTensor):
        inputs = torch.cat((z_vars[masked_var_idx], h_vars[masked_var_idx]), dim=-1)
        cons_pred = self.var_type_predictor(inputs)
        return self.convert(cons_pred)

class Bias_Predictor(nn.Module):
    def __init__(self, config, data_stats):
        """
        Bias predictor.
        """
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.rhs_type = data_stats["rhs_type"]

        self.rhs_min = int(data_stats["rhs_min"])
        self.rhs_max = int(data_stats["rhs_max"])

        self.cons_predictor = nn.Sequential(
            nn.Linear(2 * self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1),
            nn.Sigmoid(),
        )

        self.cons_loss = MSELoss(reduction="sum")
    
    def convert(self, cons: torch.Tensor):
        if self.rhs_type == "int":
            return torch.round(cons * (self.rhs_max - self.rhs_min) + self.rhs_min)
        elif self.rhs_type == "float":
            return cons * (self.rhs_max - self.rhs_min) + self.rhs_min

    def forward(self,
            z_conss: torch.Tensor,
            h_conss: torch.Tensor,
            masked_cons_idx: torch.LongTensor,
            cons_label: torch.Tensor
        ):
        inputs = torch.cat((z_conss[masked_cons_idx], h_conss[masked_cons_idx]), dim=-1)
        cons_pred = self.cons_predictor(inputs).view(-1)
        if abs(self.rhs_max - self.rhs_min) < 1e-6:
            cons_loss = torch.tensor(0.0, device=inputs.device)
        else:
            cons_label_ = (cons_label.view(-1) - self.rhs_min) / (self.rhs_max - self.rhs_min)
            cons_loss = self.cons_loss(cons_pred, cons_label_)

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("Constraint_predictor/cons_loss", cons_loss, tb_writter.step)
            tb.add_histogram("Constraint_predictor/cons", cons_label, tb_writter.step)
            tb.add_histogram("Constraint_predictor/cons_pred", self.convert(cons_pred), tb_writter.step)
        else:
            tb.add_scalar("Constraint_predictor/cons_loss_val", cons_loss, tb_writter.step)
            tb.add_histogram("Constraint_predictor/cons_val", cons_label, tb_writter.step)
            tb.add_histogram("Constraint_predictor/cons_pred_val", self.convert(cons_pred), tb_writter.step)

        return cons_loss, cons_pred

    def decode(self,  z_conss: torch.Tensor, h_conss: torch.Tensor, masked_cons_idx: torch.LongTensor):
        inputs = torch.cat((z_conss[masked_cons_idx], h_conss[masked_cons_idx]), dim=-1)
        cons_pred = self.cons_predictor(inputs).view(-1)
        return self.convert(cons_pred)

class Degree_Predictor(nn.Module):

    def __init__(self, config, data_stats):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.degree_min = data_stats["cons_degree_min"]
        self.degree_max = data_stats["cons_degree_max"]

        self.degree_predictor =  nn.Sequential(
            nn.Linear(2 * self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1),
            nn.Sigmoid()
        )

        self.degree_loss = MSELoss(reduction="sum")
        
    def convert(self, degree: torch.Tensor):
        return torch.round(degree * (self.degree_max - self.degree_min) + self.degree_min)
    
    def forward(self, z_conss: torch.Tensor, h_conss: torch.Tensor, masked_cons_idx: torch.LongTensor, degree_label: torch.Tensor):
        inputs = torch.cat((z_conss[masked_cons_idx], h_conss[masked_cons_idx]), dim=-1)
        degree_pred = self.degree_predictor(inputs).view(-1)
        degree_label_ = (degree_label - self.degree_min) / (self.degree_max - self.degree_min)
        degree_loss = self.degree_loss(degree_pred, degree_label_)

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("Degree_predictor/degree_loss", degree_loss, tb_writter.step)
            tb.add_histogram("Degree_predictor/degree", degree_label, tb_writter.step)
            tb.add_histogram("Degree_predictor/degree_pred", self.convert(degree_pred), tb_writter.step)
            tb.add_histogram("Embeddings/h_masked_cons", h_conss[masked_cons_idx], tb_writter.step)
        else:
            tb.add_scalar("Degree_predictor/degree_loss_val", degree_loss, tb_writter.step)
            tb.add_histogram("Degree_predictor/degree_val", degree_label, tb_writter.step)
            tb.add_histogram("Degree_predictor/degree_pred_val", self.convert(degree_pred), tb_writter.step)
            tb.add_histogram("Embeddings/h_masked_cons_val", h_conss[masked_cons_idx], tb_writter.step)

        return degree_loss, degree_pred
    
    def decode(self, z_conss: torch.Tensor, h_conss: torch.Tensor, masked_cons_idx: torch.LongTensor):
        inputs = torch.cat((z_conss[masked_cons_idx], h_conss[masked_cons_idx]), dim=-1)
        degree_pred = self.degree_predictor(inputs).view(-1)
        return self.convert(degree_pred)

class Logits_Predictor(nn.Module):

    def __init__(self, config, data_stats):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.logits_predictor = nn.Sequential(
            nn.Linear(2 * self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1)
        )

        self.logits_loss = BCEWithLogitsLoss(reduction='sum')

    def forward(self, z_vars: torch.Tensor, h_vars: torch.Tensor, logits_label: torch.Tensor):
       
        logits_input = torch.cat([z_vars, h_vars], dim=-1)
        logits_input, logits_label = downsample(logits_input, logits_label)
        logits_pred = self.logits_predictor(logits_input).view(-1)
        logits_loss = self.logits_loss(logits_pred, logits_label.float())

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("Logits_predictor/logits_loss", logits_loss, tb_writter.step)

            logits_label = logits_label.cpu().numpy()
            logits_pred = logits_pred.detach().cpu().numpy()
            auc = roc_auc_score(logits_label, logits_pred)
            tb.add_scalar("Logits_predictor/logits_auc", auc, tb_writter.step)
            logits_pred = (logits_pred > 0)
            tb.add_histogram("Logits_predictor/logits", logits_label, tb_writter.step)
            tb.add_histogram("Logits_predictor/logits_pred", logits_pred, tb_writter.step)

        else:
            tb.add_scalar("Logits_predictor/logits_loss_val", logits_loss, tb_writter.step)

            logits_label = logits_label.cpu().numpy()
            logits_pred = logits_pred.detach().cpu().numpy()
            auc = roc_auc_score(logits_label, logits_pred)
            tb.add_scalar("Logits_predictor/logits_auc_val", auc, tb_writter.step)
            logits_pred = (logits_pred > 0)
            tb.add_histogram("Logits_predictor/logits_val", logits_label, tb_writter.step)
            tb.add_histogram("Logits_predictor/logits_pred_val", logits_pred, tb_writter.step)

        return logits_loss, logits_pred
        
    
    def decode(self, z_vars: torch.Tensor, h_vars: torch.Tensor, x_variables_batch: torch.Tensor, degree: torch.LongTensor):
        logits_input = torch.cat([z_vars, h_vars], dim=-1)
        logits_pred = self.logits_predictor(logits_input).view(-1)
        logits_pred = unbatch(logits_pred, x_variables_batch)

        logits = [logits_pred[i].view(-1) for i in range(len(logits_pred))]

        non_zeors = []
        for degree_, logits_ in zip(degree, logits):
            degree__ = torch.minimum(degree_, torch.tensor(logits_.shape[0]))
            _, indices = torch.topk(logits_, int(degree__.item()))
            non_zeors.append(indices)
            
        return non_zeors

class Weights_Predictor(nn.Module):

    def __init__(self, config, data_stats):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.embd_size = config.embd_size

        self.weights_type = data_stats["coef_type"]
        self.weights_min = int(data_stats["coef_min"])
        self.weights_max = int(data_stats["coef_max"])

        self.weights_predictor = nn.Sequential(
            nn.Linear(2 * self.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, 1),
            nn.Sigmoid()
        )

        self.weights_loss = MSELoss(reduction="sum")

    def convert(self, weights: torch.Tensor):
        if self.weights_type == "int":
            return torch.round(weights * (self.weights_max - self.weights_min) + self.weights_min)
        else:
            return weights * (self.weights_max - self.weights_min) + self.weights_min
        
    def forward(self, z_vars, h_vars, weights_idx, weights_label):
        weights_input = torch.cat((z_vars[weights_idx], h_vars[weights_idx]), dim=-1)
        weights_pred = self.weights_predictor(weights_input).view(-1)
        if abs(self.weights_max - self.weights_min) < 1e-6:
            weights_loss = torch.tensor(0.0).to(weights_pred.device)
        else:
            weights_label_ = (weights_label.view(-1) - self.weights_min) / (self.weights_max - self.weights_min)
            weights_loss = self.weights_loss(weights_pred, weights_label_)

        tb = tb_writter.tb_writter
        if self.training:
            tb.add_scalar("Weights_predictor/weights_loss", weights_loss, tb_writter.step)
            tb.add_histogram("Weights_predictor/weights", weights_label, tb_writter.step)
            tb.add_histogram("Weights_predictor/weights_pred", self.convert(weights_pred), tb_writter.step)
        else:
            tb.add_scalar("Weights_predictor/weights_loss_val", weights_loss, tb_writter.step)
            tb.add_histogram("Weights_predictor/weights_val", weights_label, tb_writter.step)
            tb.add_histogram("Weights_predictor/weights_pred_val", self.convert(weights_pred), tb_writter.step)
        return weights_loss, weights_pred

    def decode(self, z_vars, h_vars, x_variables_batch, weights_idx):
        weights_input = []
        for h_vars_, z_, weights_idx_ in zip(unbatch(h_vars, x_variables_batch), unbatch(z_vars, x_variables_batch), weights_idx):
            weights_input.append(torch.cat((h_vars_[weights_idx_], z_[weights_idx_]), dim=-1))
        weights_input = torch.cat(weights_input, dim=0)
        weights = self.weights_predictor(weights_input)
        return self.convert(weights)
src/modules.py
MIRALab-USTC/L2O-G2MILP
from typing import Tuple, Union

import torch
import torch.nn as nn
from omegaconf import DictConfig
from torch import Tensor
from torch_geometric.data import Batch
from torch_geometric.nn import (GraphNorm, JumpingKnowledge, MessagePassing,
                                aggr)
from torch_geometric.utils import scatter

from src.data import BipartiteGraph


class BipartiteGraphEmbedding(nn.Module):
    def __init__(self, config: DictConfig) -> None:
        """
        Embed the bipartite graph into the embedding space.
        """
        super().__init__()
        self.emb_size = config.embd_size
        self.hidden_size = config.hidden_size
        self.cons_nfeats = 1
        self.edge_nfeats = 1
        self.var_nfeats = 9

        self.cons_embedding = nn.Sequential(
            nn.Linear(self.cons_nfeats, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.edge_embedding = nn.Sequential(
            nn.Linear(self.edge_nfeats, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.var_embedding = nn.Sequential(
            nn.Linear(self.var_nfeats, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.mask_embedding = nn.Parameter(
            torch.randn(self.emb_size), requires_grad=True)

    def embed_graph(self, graph: Union[Batch, BipartiteGraph]) -> Union[Batch, BipartiteGraph]:
        """
        Embed the bipartite graphs into the embedding space.

        Input:
            graph: a batch of bipartite graphs

        Output:
            graph: a batch of bipartite graphs with embedded features
        """
        graph = graph.clone()
        graph.x_constraints = self.cons_embedding.forward(graph.x_constraints)
        graph.edge_attr = self.edge_embedding.forward(graph.edge_attr)
        graph.x_variables = self.var_embedding.forward(graph.x_variables)
        return graph

    def mask_and_embed_graph(self, graph: Union[Batch, BipartiteGraph]) -> Union[Batch, BipartiteGraph]:
        """
        Randomly mask a constraint node of each input graph. Then embed the masked bipartite graphs into the embedding space.

        Input:
            graph: a batch of bipartite graphs

        Output:
            graph: a batch of masked bipartite graphs with embedded features
        """
        if isinstance(graph, list):
            graph_list = graph.copy()
        else:
            graph_list = graph.to_data_list()

        for bid, graph in enumerate(graph_list):
            # get the masked constraint idx
            masked_cons_idx = torch.randint(
                0, graph.num_constraints, (1,)).item()

            # get the labels
            masked_edges_idx = torch.where(
                graph.edge_index[0] == masked_cons_idx)[0]
            connected_vars_idx = graph.edge_index[1][masked_edges_idx]
            logits_label = torch.zeros(graph.num_variables, dtype=int)
            logits_label[connected_vars_idx] = 1
            degree_label = torch.sum(logits_label)

            device = graph.x_constraints.device
            graph.masked_cons_idx = masked_cons_idx
            graph.bias_label = graph.x_constraints[masked_cons_idx].to(device)
            graph.degree_label = degree_label.view(1).to(device)
            graph.logits_label = logits_label.to(device)
            graph.connected_vars_idx = connected_vars_idx.to(device)
            graph.weights_label = graph.edge_attr[masked_edges_idx].to(device)

            # update the graph
            graph = self.embed_graph(graph)
            seen_edges_idx = torch.where(
                graph.edge_index[0] != masked_cons_idx)[0]
            graph.edge_index = graph.edge_index[:, seen_edges_idx]
            graph.edge_attr = graph.edge_attr[seen_edges_idx]
            graph.x_constraints[masked_cons_idx] = self.mask_embedding

            graph_list[bid] = graph

        return Batch.from_data_list(graph_list, follow_batch=["x_constraints", "x_variables"])


class BipartiteGraphGNN(nn.Module):
    def __init__(self,
                 config: DictConfig,
                 is_masked: bool = False
                 ):
        """
        The GNN model for bipartite graphs.

        Args:
            config: the configuration of the GNN model
            is_masked: whether the input graph is masked
        """
        super().__init__()
        self.emb_size = config.embd_size
        self.hidden_size = config.hidden_size
        self.depth = config.depth
        self.cons_nfeats = 1
        self.edge_nfeats = 1
        self.var_nfeats = 9
        self.jk = config.jk
        self.is_masked = is_masked

        self.conv_v_to_c_layers = nn.Sequential()
        self.conv_c_to_v_layers = nn.Sequential()
        self.graph_norm_v_to_c_layers = nn.Sequential()
        self.graph_norm_c_to_v_layers = nn.Sequential()

        for _ in range(self.depth):
            self.conv_v_to_c_layers.append(BipartiteGraphConvolution(
                config, aggr_coef=config.aggr_coef.v_to_c))
            self.graph_norm_v_to_c_layers.append(GraphNorm(self.emb_size))
            self.conv_c_to_v_layers.append(BipartiteGraphConvolution(
                config, aggr_coef=config.aggr_coef.c_to_v))
            self.graph_norm_c_to_v_layers.append(GraphNorm(self.emb_size))

        if self.is_masked:
            self.virtual_c_to_v_layers = nn.Sequential()
            self.virtual_v_to_c_layers = nn.Sequential()
            for _ in range(self.depth):
                self.virtual_v_to_c_layers.append(VirtualAggr(config))
                self.virtual_c_to_v_layers.append(nn.Sequential(
                    nn.Linear(2 * self.emb_size, self.hidden_size),
                    nn.ReLU(),
                    nn.Linear(self.hidden_size, self.emb_size)
                ))

        if self.jk == 'cat':
            self.var_jk_layer = nn.Sequential(
                JumpingKnowledge(
                    mode=self.jk, channels=self.emb_size, num_layers=self.depth + 1),
                nn.Linear(self.emb_size * (self.depth + 1), self.hidden_size),
                nn.ReLU(),
                nn.Linear(self.hidden_size, self.emb_size)
            )
            self.cons_jk_layer = nn.Sequential(
                JumpingKnowledge(
                    mode=self.jk, channels=self.emb_size, num_layers=self.depth + 1),
                nn.Linear(self.emb_size * (self.depth + 1), self.hidden_size),
                nn.ReLU(),
                nn.Linear(self.hidden_size, self.emb_size)
            )
        elif self.jk == 'lstm':
            self.var_jk_layer = JumpingKnowledge(
                mode=self.jk, channels=self.emb_size, num_layers=self.depth + 1)
            self.cons_jk_layer = JumpingKnowledge(
                mode=self.jk, channels=self.emb_size, num_layers=self.depth + 1)

    def forward(self, graph: BipartiteGraph) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Forward pass of the GNN model.

        Input:
            graph: a batch of bipartite graphs

        Output:
            h_conss: the final embedding of constraint nodes
            h_vars: the final embedding of variable nodes
        """
        constraint_features = graph.x_constraints
        edge_indices = graph.edge_index
        edge_features = graph.edge_attr
        variable_features = graph.x_variables

        constraint_features_batch = graph.x_constraints_batch
        variable_features_batch = graph.x_variables_batch

        reversed_edge_indices = torch.stack(
            [edge_indices[1], edge_indices[0]], dim=0)

        if self.jk:
            constraint_features_list, variable_features_list = [], []
            constraint_features_list.append(constraint_features)
            variable_features_list.append(variable_features)

        for i in range(self.depth):
            constraint_features = self.conv_v_to_c_layers[i](
                variable_features, reversed_edge_indices, edge_features, constraint_features
            )
            variable_features = self.conv_c_to_v_layers[i](
                constraint_features, edge_indices, edge_features, variable_features
            )
            if self.is_masked and hasattr(graph, "masked_cons_idx"):
                virtual_aggr_info = self.virtual_v_to_c_layers[i](
                    variable_features, constraint_features[graph.masked_cons_idx], graph.x_variables_batch
                )
                constraint_features[graph.masked_cons_idx] = virtual_aggr_info
                variable_features = self.virtual_c_to_v_layers[i](
                    torch.cat(
                        (variable_features, virtual_aggr_info[graph.x_variables_batch]), dim=-1)
                )

            constraint_features = self.graph_norm_v_to_c_layers[i](
                constraint_features, constraint_features_batch)
            variable_features = self.graph_norm_c_to_v_layers[i](
                variable_features, variable_features_batch)

            if self.jk:
                constraint_features_list.append(constraint_features)
                variable_features_list.append(variable_features)

        if self.jk:
            h_conss = self.cons_jk_layer(constraint_features_list)
            h_vars = self.var_jk_layer(variable_features_list)
        else:
            h_conss = constraint_features
            h_vars = variable_features

        return h_conss, h_vars


class BipartiteGraphConvolution(MessagePassing):
    def __init__(self, config, aggr_coef=None):
        if config.aggr in ["add", "mean", "max", "min"]:
            super().__init__(aggr=config.aggr)
        elif config.aggr == "softmax":
            super().__init__(aggr=aggr.SoftmaxAggregation(learn=True))
        else:
            raise NotImplementedError

        self.emb_size = config.embd_size
        self.hidden_size = config.hidden_size
        self.aggr_coef = aggr_coef

        self.feature_module_left = nn.Sequential(
            nn.Linear(self.emb_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.feature_module_edge = nn.Sequential(
            nn.Linear(self.emb_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.feature_module_right = nn.Sequential(
            nn.Linear(self.emb_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.feature_module_final = nn.Sequential(
            nn.Linear(self.emb_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

        self.output_module = nn.Sequential(
            nn.Linear(2 * self.emb_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, self.emb_size)
        )

    def forward(self, left_features, edge_indices, edge_features, right_features):
        output = self.propagate(
            edge_indices,
            size=(left_features.shape[0], right_features.shape[0]),
            node_features=(left_features, right_features),
            edge_features=edge_features,
        )
        return self.output_module(
            torch.cat([output, right_features], dim=-1)
        )

    def message(self, node_features_i, node_features_j, edge_features):
        output = self.feature_module_final(
            self.feature_module_left(node_features_i)
            + self.feature_module_edge(edge_features)
            + self.feature_module_right(node_features_j)
        )
        return output


class VirtualAggr(torch.nn.Module):

    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.output_module = nn.Sequential(
            nn.Linear(2 * config.embd_size, self.hidden_size),
            nn.ReLU(),
            nn.Linear(self.hidden_size, config.embd_size)
        )

    def forward(
        self,
        x_variables: Tensor,
        x_virtual_constraints: Tensor,
        x_variables_batch: Tensor,
    ):
        x = scatter(x_variables, x_variables_batch, reduce='mean', dim=0)
        return self.output_module(
            torch.cat([x, x_virtual_constraints], dim=-1)
        )

src/nn.py
MIRALab-USTC/L2O-G2MILP
from tensorboardX import SummaryWriter

tb_writter: SummaryWriter = None
step: int = 0

def set_logger(log_dir: str):
    global tb_writter
    global step
    tb_writter = SummaryWriter(log_dir)
    step = 0

def set_step(total_step: int):
    global step
    step = total_step

def add_scalar(tag, scalar_value, global_step=None, walltime=None):
    global tb_writter
    tb_writter.add_scalar(tag, scalar_value, global_step, walltime)
src/tb_writter.py
MIRALab-USTC/L2O-G2MILP
import logging
import math
import os
import os.path as path
from typing import Iterator, Union

import torch
import torch.optim as optim
from omegaconf import DictConfig
from torch_geometric.data import Batch
from torch_geometric.loader import DataLoader

import src.tb_writter as tb_writter
from src.benchmarks.benchmarks import Benchmark
from src.data import BipartiteGraph, InstanceDataset, TrainSampler
from src.generator import Generator
from src.model import G2MILP


class Trainer():
    def __init__(self,
                 model: G2MILP,
                 train_set: InstanceDataset,
                 paths: DictConfig,
                 config: DictConfig,
                 generator_config: DictConfig,
                 benchmark_config: DictConfig,
                 ):
        """
        Trainer for G2MILP.

        Args:
            model: G2MILP model
            train_set: training dataset
            eval_set: dev dataset
            paths: paths
            config: trainer config
            generator_config: generator config
            benchmark_config: benchmark config
        """
        self.model = model

        # laod paths
        self.paths = paths
        self.model_dir = paths.model_dir
        self.samples_dir = paths.samples_dir
        self.dataset_samples_dir = paths.dataset_samples_dir
        self.dataset_stats_dir = paths.dataset_stats_dir
        self.benchmark_dir = paths.benchmark_dir
        os.makedirs(self.model_dir, exist_ok=True)

        # load parameters
        self.total_steps = config.steps
        self.step = 0
        self.save_step = config.save_step

        self.grad_max_norm = config.grad_max_norm

        # prepare the data loader
        self.batch_size = config.batch_size
        self.batch_repeat_size = config.batch_repeat_size

        self.train_loader: Iterator[BipartiteGraph] = DataLoader(
            dataset=train_set,
            batch_sampler=TrainSampler(
                train_set,
                batch_size=self.batch_size,
                repeat_size=self.batch_repeat_size,
                total_steps=self.total_steps,
            ),
            follow_batch=["x_constraints", "x_variables"],
        )
    
        # init optimizer
        self.optimizer = optim.Adam(
            model.parameters(), lr=config.lr.init, weight_decay=config.weight_decay)

        # init lr scheduler
        self.lr_anneal_step = config.lr.anneal_step
        self.lr_scheduler = optim.lr_scheduler.ExponentialLR(
            self.optimizer, config.lr.anneal_rate)

        # init beta scheduler
        self.beta_cons, self.beta_var = config.beta.cons.min, config.beta.var.min
        self.beta_cons_scheduler = beta_scheduler(config.beta.cons)
        self.beta_var_scheduler = beta_scheduler(config.beta.var)

        self.performance_track = []
        self.save_start = config.save_start

        self.generator_config = generator_config
        self.benchmark_config = benchmark_config

    def train(self):
        """
        Train the model.
        """
        for data in self.train_loader:
            data = data.cuda()
            self.step += 1
            tb_writter.set_step(self.step)
            self.step_train(data)

            self.beta_cons = self.beta_cons_scheduler.step()
            self.beta_var = self.beta_var_scheduler.step()

            if self.step % self.save_step == 0 and self.step >= self.save_start:
                self.step_save()

            if self.step % self.lr_anneal_step == 0:
                self.lr_scheduler.step()

        self.save_best_ckpt()

    def step_train(self, data: Union[Batch, BipartiteGraph]):
        """
        Train the model for one step.
        """
        self.model.train()
        self.model.zero_grad()
        torch.cuda.empty_cache()
        loss = self.model.forward(
            data, beta_cons=self.beta_cons, beta_var=self.beta_var)
        loss.backward()

        tb_writter.add_scalar("Train/lr", self.lr, self.step)
        tb_writter.add_scalar("Train/beta_cons", self.beta_cons, self.step)
        tb_writter.add_scalar("Train/beta_var", self.beta_var, self.step)
        tb_writter.add_scalar("Train/total_loss", loss, self.step)

        torch.nn.utils.clip_grad_norm_(
            self.model.parameters(), self.grad_max_norm)
        self.optimizer.step()

        logging.info(
            f"Step {self.step}/{self.total_steps}. Loss: {loss.item():.5f}. beta_cons: {self.beta_cons:.5f}. beta_var: {self.beta_var:.5f}. lr: {self.lr}.")

    def step_save(self):
        """
        Save the model and generate samples.
        """
        self.model.eval()
        self.model.zero_grad()
        torch.cuda.empty_cache()
        save_path = path.join(self.model_dir, f"model_step_{self.step}.ckpt")
        with torch.no_grad():
            torch.save(self.model.state_dict(), save_path)

        samples_dir = self.samples_dir + f"_step_{self.step}"
        os.makedirs(samples_dir, exist_ok=True)
        generator = Generator(
            model=self.model,
            config=self.generator_config,
            templates_dir=self.dataset_samples_dir,
            save_dir=samples_dir
        )
        generator.generate()

        benchmark_dir = self.benchmark_dir + f"_step_{self.step}"
        benchmarker = Benchmark(
            config=self.benchmark_config,
            dataset_stats_dir=self.dataset_stats_dir,
        )
        results = benchmarker.assess_samples(
            samples_dir=samples_dir,
            benchmark_dir=benchmark_dir
        )

        info_path = path.join(benchmark_dir, "info.json")
        benchmarker.log_info(
            generator_config=self.generator_config,
            benchmarking_config=self.benchmark_config,
            meta_results=results,
            save_path=info_path,
        )

        tb_writter.add_scalar("Train/similarity",
                              results["distribution"]["score"], self.step)
        tb_writter.add_scalar("Train/solving_time_mean",
                              results["solving"]["solving_time"]["mean"], self.step)
        tb_writter.add_scalar(
            "Train/solving_time_err", results["solving"]["solving_time"]["mean_error"], self.step)
        tb_writter.add_scalar("Train/num_nodes_mean",
                              results["solving"]["num_nodes"]["mean"], self.step)
        tb_writter.add_scalar(
            "Train/num_nodes_err", results["solving"]["num_nodes"]["mean_error"], self.step)

        # we use the sum of mean error as the metric to find the best model
        err = results["solving"]["solving_time"]["mean_error"] + \
            results["solving"]["num_nodes"]["mean_error"]
        self.performance_track.append((self.step, err))

    def save_best_ckpt(self):
        """
        Save the best model.
        """
        self.performance_track.sort(key=lambda x: x[1])
        best_step = self.performance_track[0][0]

        logging.info(f"Best step: {best_step}.")
        best_model_path = path.join(
            self.model_dir, f"model_step_{best_step}.ckpt")
        self.model.eval()
        self.model.zero_grad()
        self.model.load_state_dict(torch.load(best_model_path))
        torch.cuda.empty_cache()
        save_path = path.join(self.model_dir, f"model_best.ckpt")
        with torch.no_grad():
            torch.save(self.model.state_dict(), save_path)
        logging.info(f"Best model saved to {save_path}.")

    @property
    def lr(self):
        return self.lr_scheduler.get_last_lr()[0]


class scheduler(object):
    def __init__(self, config) -> None:
        self.beta = config.min
        self.t = 0
        self.warmup = config.warmup
        self.beta_min = config.min
        self.beta_max = config.max
        self.beta_anneal_period = config.anneal_period

    def step(self):
        pass


class linear_schedule(scheduler):
    def __init__(self, config) -> None:
        super().__init__(config)
        self.beta_step = (self.beta_max - self.beta_min) / \
            self.beta_anneal_period

    def step(self):
        if self.t < self.warmup:
            self.beta = self.beta_min
        elif self.t <= self.warmup + self.beta_anneal_period:
            self.beta += self.beta_step
        self.t += 1
        return self.beta


class cyclical_schedule(scheduler):
    def __init__(self, config) -> None:
        super().__init__(config)
        self.beta_num_cycles = config.num_cycles

        self.cycle_period = self.beta_anneal_period // self.beta_num_cycles
        self.linear_period = int(self.cycle_period * 0.5)
        self.beta_step = (self.beta_max - self.beta_min) / self.linear_period

        self.T = max((self.t - self.warmup) // self.cycle_period + 1, 1)
        self.tau = self.t - self.warmup - self.T * self.cycle_period

    def step(self):
        if self.t < self.warmup:
            self.beta = self.beta_min
            self.t += 1
        else:
            if self.tau == 0 and self.T < self.beta_num_cycles:
                self.beta = self.beta_min
                self.T += 1
            elif self.tau <= self.linear_period:
                self.beta = min(self.beta + self.beta_step, self.beta_max)
            self.tau = (self.tau + 1) % self.cycle_period
            self.t += 1

        return self.beta


class sigmoid_schedule(scheduler):
    def __init__(self, config) -> None:
        super(sigmoid_schedule, self).__init__(config)
        self.diff = self.beta_max - self.beta_min
        self.anneal_rate = math.pow(0.01, 1 / self.beta_anneal_period)
        self.weight = 1

    def step(self):
        if self.t < self.warmup:
            self.beta = self.beta_min
        else:
            self.weight = math.pow(self.anneal_rate, self.t - self.warmup)
            self.beta = self.beta_min + self.diff * (1 - self.weight)
        self.t += 1
        return self.beta


class beta_scheduler(object):
    def __init__(self, config) -> None:
        self.mode = config.mode
        if self.mode not in ["linear", "sigmoid", "cyclical"]:
            self.mode = "sigmoid"
        if self.mode == "linear":
            self.schedule = linear_schedule(config)
        elif self.mode == "cyclical":
            self.schedule = cyclical_schedule(config)
        else:
            self.schedule = sigmoid_schedule(config)

    def step(self):
        return self.schedule.step()

src/trainer.py
MIRALab-USTC/L2O-G2MILP

import os
import random

import community.community_louvain as community
import ecole
import gurobipy as gp
import networkx as nx
import numpy as np
import pyscipopt as scip
import torch
from scipy.sparse import coo_matrix
from torch_geometric.data import Data
from torch_geometric.utils import to_networkx

from src.data import BipartiteGraph

COEF_TYPES = {
    'debug': 'int',
    'setcover': 'int',
    'mis': 'int',
}

RHS_TYPES = {
    'debug': 'int',
    'setcover': 'int',
    'mis': 'int',
}

VAR_FEATURES = ["objective", "is_type_binary", "is_type_integer", "is_type_implicit_integer", "is_type_continuous", "has_lower_bound", "has_upper_bound", "lower_bound", "upper_bound"]

VAR_TYPE_FEATURES = ["B", "I", "M", "C"]

def instance2graph(path: str, compute_features: bool = False):
    """
    Extract the bipartite graph from the instance file.
    compute_features: whether to compute the features of the instance
    """
    model = ecole.scip.Model.from_file(path)
    obs = ecole.observation.MilpBipartite().extract(model, True)
    constraint_features = obs.constraint_features
    edge_indices = np.array(obs.edge_features.indices, dtype=int)
    edge_features = obs.edge_features.values.reshape((-1,1))
    variable_features = obs.variable_features
    graph = [constraint_features, edge_indices, edge_features, variable_features]
    if not compute_features:
        return graph, None
    else:
        n_conss = len(constraint_features)
        n_vars = len(variable_features)
        n_cont_vars = np.sum(variable_features, axis=0)[1]
        
        lhs = coo_matrix((edge_features.reshape(-1), edge_indices), shape=(n_conss, n_vars)).toarray()
        rhs = constraint_features.flatten()
        obj = variable_features[:, 0].flatten()

        nonzeros = (lhs != 0)
        n_nonzeros = np.sum(nonzeros)
        lhs_coefs = lhs[np.where(nonzeros)]
        var_degree, cons_degree = nonzeros.sum(axis=0), nonzeros.sum(axis=1)

        pyg_graph = Data(
            x_s = constraint_features,
            x_t = variable_features,
            edge_index = torch.LongTensor(edge_indices),
            node_attribute = "bipartite"
        )
        pyg_graph.num_nodes = len(constraint_features) + len(variable_features)
        nx_graph = to_networkx(pyg_graph, to_undirected=True)

        features = {
            "instance": path,
 
            "n_conss": n_conss,
            "n_vars": n_vars,
            "n_cont_vars": n_cont_vars,
            "ratio_cont_vars": float(n_cont_vars / n_vars),

            "n_nonzeros": n_nonzeros,
            "coef_dens": float(len(edge_features) / (n_vars * n_conss)),

            "var_degree_mean": float(var_degree.mean()),
            "var_degree_std": float(var_degree.std()),
            "var_degree_min": float(var_degree.min()),
            "var_degree_max": float(var_degree.max()),

            "cons_degree_mean": float(cons_degree.mean()),
            "cons_degree_std": float(cons_degree.std()),
            "cons_degree_min": int(cons_degree.min()),
            "cons_degree_max": int(cons_degree.max()),

            "lhs_mean": float(lhs_coefs.mean()),
            "lhs_std": float(lhs_coefs.std()),
            "lhs_min": float(lhs_coefs.min()),
            "lhs_max": float(lhs_coefs.max()),

            "rhs_mean": float(rhs.mean()),
            "rhs_std": float(rhs.std()),
            "rhs_min": float(rhs.min()),
            "rhs_max": float(rhs.max()),

            "obj_mean": float(obj.mean()),
            "obj_std": float(obj.std()),
            "obj_min": float(obj.min()),
            "obj_max": float(obj.max()),

            "clustering": float(nx.average_clustering(nx_graph)),
            "modularity": float(community.modularity(community.best_partition(nx_graph), nx_graph)),
        }
        return graph, features

def solve_instance(path: str):
    """
    Solve the instance using Gurobi.
    """
    env = gp.Env(empty=True)
    env.setParam("OutputFlag", 0)
    env.setParam("TimeLimit", 300)
    env.setParam("Threads", 1)
    env.start()
    model = gp.read(path, env=env)
    model.optimize()

    results = {
        "status": model.status,
        "obj": model.objVal,
        "num_nodes": model.NodeCount,
        "num_sols": model.SolCount,
        "solving_time": model.Runtime,
    }
    return results


def mask_graph(graph: BipartiteGraph):
    """
    Randomly mask a constraint node of the input graph.
    """
    masked_cons_idx = torch.randint(0, len(graph.x_constraints), (1,)).item()
    graph.masked_cons_idx = masked_cons_idx

    # update the graph
    seen_edges = torch.where(graph.edge_index[0] != masked_cons_idx)[0]
    graph.edge_index = graph.edge_index[:, seen_edges]
    graph.edge_attr = graph.edge_attr[seen_edges]
    graph.edge_index[0, graph.edge_index[0] > masked_cons_idx] -= 1
    seen_constraints = torch.tensor(list(range(masked_cons_idx)) + list(range(masked_cons_idx+1, len(graph.x_constraints))))
    graph.x_constraints = graph.x_constraints[seen_constraints]
    graph.num_nodes = len(graph.x_constraints) + len(graph.x_variables)
    return graph

def model_add_var(model: scip.Model, var_idx: int, var_features: np.ndarray):
    """
    Add a variable to the SCIP model.
    model: the SCIP model.
    var_idx: the index of the added variable.
    var_features: the features of the variable. Form: [objective, is_type_binary, is_type_integer, is_type_implicit_integer, is_type_continuous, has_lower_bound, has_upper_bound, lower_bound, upper_bound].
    """
    return model.addVar(
        name = f"x{var_idx}",
        obj = var_features[0],
        vtype = VAR_TYPE_FEATURES[np.argmax(var_features[1:5])],
        lb = var_features[7] if bool(var_features[5]) else None,
        ub = var_features[8] if bool(var_features[6]) else None,
    )

def graph2instance(graph):
    [constraint_features, edge_indices, edge_features, variable_features] = graph
    
    model = scip.Model()
    model.setIntParam('display/verblevel', 0)

    vars = [model_add_var(model, i, var_features) for i, var_features in enumerate(variable_features)]
    for cons_idx, cons in enumerate(constraint_features):
        edge_idx = np.where(edge_indices[0] == cons_idx)[0]
        cons_expr = sum([edge_features[i,0] * vars[edge_indices[1,i]] for i in edge_idx])
        cons_expr.normalize()
        model.addCons(cons_expr <= cons[0])

    return model


def downsample(X: torch.Tensor, y: torch.LongTensor):
    """
    Downsample the input data and target data so that the negative samples are balanced with the positive samples.
    """
    pos_idx = torch.where(y == 1)[0]
    neg_idx = torch.where(y == 0)[0]
    neg_idx = neg_idx[torch.randperm(neg_idx.shape[0])[:pos_idx.shape[0]]]
    idx = torch.cat([pos_idx, neg_idx])
    idx = idx[torch.randperm(idx.shape[0])]
    return X[idx], y[idx]


def set_seed(seed):
    """
    Set the random seed.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    # Set CuDNN to be deterministic. Notice that this may slow down the training.
    os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'
    torch.use_deterministic_algorithms(True)
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.benchmark = True
    torch.backends.cudnn.deterministic = True

def set_cpu_num(cpu_num):
    """
    Set the number of used cpu kernals.
    """
    os.environ['OMP_NUM_THREADS'] = str(cpu_num)
    os.environ['OPENBLAS_NUM_THREADS'] = str(cpu_num)
    os.environ['MKL_NUM_THREADS'] = str(cpu_num)
    os.environ['VECLIB_MAXIMUM_THREADS'] = str(cpu_num)
    os.environ['NUMEXPR_NUM_THREADS'] = str(cpu_num)
    torch.set_num_threads(cpu_num)
src/utils.py
MIRALab-USTC/L2O-G2MILP
import logging
import os.path as path

import hydra
import torch
from omegaconf import DictConfig

import src.tb_writter as tb_writter
from src import (G2MILP, Benchmark, Generator, InstanceDataset, Trainer,
                 set_cpu_num, set_seed)


@hydra.main(version_base=None, config_path="conf", config_name="train")
def train(config: DictConfig):
    """
    Train G2MILP.
    """
    set_seed(config.seed)
    set_cpu_num(config.num_workers + 1)
    torch.cuda.set_device(config.cuda)
    tb_writter.set_logger(config.paths.tensorboard_dir)

    model = G2MILP.load_model(config)
    logging.info(f"Loaded model.")
    logging.info(
        f"  Number of model parameters: {sum([x.nelement() for x in model.parameters()]) / 1000}K.")

    train_set = InstanceDataset(
        data_dir=config.paths.dataset_samples_dir,
        solving_results_path=config.paths.dataset_solving_path,
    )
    logging.info(f"Loaded dataset.")
    logging.info(f"  Number of training instances: {len(train_set)}.")
   
    trainer = Trainer(
        model=model,
        train_set=train_set,
        paths=config.paths,
        config=config.trainer,
        generator_config=config.generator,
        benchmark_config=config.benchmarking,
    )

    logging.info("="*10 + "Begin training" + "="*10)

    trainer.train()

    logging.info("="*10 + "Training finished" + "="*10)

    # test
    for mask_ratio in [0.01, 0.05, 0.1]:
        config.generator.mask_ratio = mask_ratio

        # generate
        samples_dir = path.join(config.paths.train_dir,
                                f"eta-{mask_ratio}/samples")
        generator = Generator(
            model=model,
            config=config.generator,
            templates_dir=config.paths.dataset_samples_dir,
            save_dir=samples_dir,
        )
        generator.generate()

        # benchmark
        benchmark_dir = path.join(
            config.paths.train_dir, f"eta-{mask_ratio}/benchmark")
        benchmarker = Benchmark(
            config=config.benchmarking,
            dataset_stats_dir=config.paths.dataset_stats_dir,
        )
        results = benchmarker.assess_samples(
            samples_dir=samples_dir,
            benchmark_dir=benchmark_dir
        )

        info_path = path.join(benchmark_dir, "info.json")
        benchmarker.log_info(
            generator_config=config.generator,
            benchmarking_config=config.benchmarking,
            meta_results=results,
            save_path=info_path,
        )


if __name__ == '__main__':
    train()

train.py
MIRALab-USTC/L2O-G2MILP