trainPredictModel.py

import argparse
import hashlib
import json
import random
import time
from pathlib import Path

import numpy as np
import torch
import torch_geometric


def EnergyWeightNorm(task):
    # Paper Eqs. (3)-(5) use exp(-c^T x) for minimization. Negative signs here
    # convert IS/CA maximization objectives to energies. Magnitudes are the
    # released repository's undocumented temperatures, not paper parameters.
    return {'IP': 1, 'WA': 100, 'IS': -100, 'CA': -1000}[task]


def train(predict, data_loader, optimizer=None, weight_norm=1, device='cpu'):
    """One epoch of the weighted cross entropy in paper Eqs. (4)-(5)."""
    predict.train(optimizer is not None)
    mean_loss, n_samples_processed = 0., 0
    with torch.set_grad_enabled(optimizer is not None):
        for batch in data_loader:
            batch = batch.to(device)
            logits = predict(batch.constraint_features, batch.edge_index,
                             batch.edge_attr, batch.variable_features)
            loss = logits.new_zeros(())
            sol_offset = val_offset = var_offset = binary_offset = 0
            for nsol, nvar, nbinary in zip(batch.nsols.tolist(), batch.ntvars.tolist(),
                                           batch.nbinary.tolist()):
                # Engineering: flat per-graph mappings batch correctly in modern
                # PyG; the released nested varInds indexing mixed graph/sample axes.
                variable_map = batch.varname_map[var_offset:var_offset + nvar].long()
                binary = batch.binary_indices[binary_offset:binary_offset + nbinary].long()
                sols = batch.solutions[sol_offset:sol_offset + nsol * nvar].reshape(nsol, nvar)
                values = batch.objVals[val_offset:val_offset + nsol]
                labels = sols[:, variable_map][:, binary]
                if not torch.isfinite(values).all() or not torch.isfinite(labels).all():
                    raise ValueError('Nonfinite labels/objectives in the training dataset.')
                if not ((labels == 0) | (labels == 1)).all():
                    raise ValueError('The binary prediction loss requires binary labels.')
                # Algebraically identical to normalized exp(-objective/temperature)
                # and BCE in Eqs. (4)-(5), with log-sum-exp numerical stabilization.
                weights = torch.softmax(-values / weight_norm, dim=0).to(logits.dtype)
                prediction = logits[var_offset:var_offset + nvar][binary]
                cross_entropy = torch.nn.functional.binary_cross_entropy_with_logits(
                    prediction.expand_as(labels), labels, reduction='none')
                loss = loss + (cross_entropy * weights[:, None]).sum()
                sol_offset += nsol * nvar
                val_offset += nsol
                var_offset += nvar
                binary_offset += nbinary
            if not torch.isfinite(loss):
                raise ValueError('Nonfinite training loss.')
            if optimizer is not None:
                optimizer.zero_grad()
                # Paper Eq. (4) and the released implementation sum instance
                # contributions. Only the displayed epoch loss is averaged.
                loss.backward()
                optimizer.step()
            mean_loss += loss.item()
            n_samples_processed += batch.num_graphs
    if n_samples_processed == 0:
        raise ValueError('Training and validation splits must both be nonempty.')
    return mean_loss / n_samples_processed


def main():
    parser = argparse.ArgumentParser(description='Train the released GNN on weighted solution pools.')
    parser.add_argument('--dataDir', default='.')
    parser.add_argument('--task', choices=['CA', 'IS', 'IP', 'WA'], default='CA')
    # Paper Sec. 4: Adam, learning rate .003, batch size 8. The paper gives no
    # stopping rule; 9999 epochs and first-50 pool solutions are released choices.
    parser.add_argument('--epochs', type=int, default=9999)
    parser.add_argument('--batch-size', type=int, default=8)
    parser.add_argument('--lr', type=float, default=0.003)
    parser.add_argument('--num-workers', type=int, default=0)
    parser.add_argument('--max-solutions', type=int, default=50)
    parser.add_argument('--temperature', type=float,
                        help='Positive energy temperature; default is the released task-specific magnitude.')
    parser.add_argument('--graph-mode', choices=['corrected', 'upstream'], default='corrected')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--device', default='auto', choices=['auto', 'cpu', 'cuda'])
    args = parser.parse_args()
    if min(args.epochs, args.batch_size, args.lr, args.max_solutions) <= 0:
        parser.error('Epochs, batch size, learning rate and solution cap must be positive.')
    if args.num_workers < 0 or args.seed < 0 or (args.temperature is not None and args.temperature <= 0):
        parser.error('Workers/seed must be nonnegative; temperature must be positive.')
    # Engineering: seed the split, model and loader. CUDA scatter reductions can
    # still be nondeterministic; CPU is the reproducibility reference.
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    torch.backends.cudnn.benchmark = False
    torch.set_num_threads(1)
    device = torch.device(('cuda' if torch.cuda.is_available() else 'cpu')
                          if args.device == 'auto' else args.device)
    root = Path(args.dataDir)
    dataset_dir = root / 'dataset' / args.task
    collection_path = dataset_dir / 'collection.json'
    if not collection_path.exists():
        raise ValueError('Run gurobi.py to collect graph/label provenance before training.')
    collection = json.loads(collection_path.read_text())
    if collection['task'] != args.task or collection['settings']['graph_mode'] != args.graph_mode:
        raise ValueError('Task/graph mode differs from collected data; use matching CLI settings.')
    names = sorted(collection['instances'])
    if set(p.name for p in (dataset_dir / 'BG').glob('*.bg')) != {name + '.bg' for name in names}:
        raise ValueError('Graph files differ from collection manifest; recollect in a clean data root.')
    sample_files = []
    for name in names:
        pair = [dataset_dir / directory / (name + suffix)
                for directory, suffix in [('BG', '.bg'), ('solution', '.sol')]]
        for path, suffix in zip(pair, ['.bg', '.sol']):
            if hashlib.sha256(path.read_bytes()).hexdigest() != collection['instances'][name]['hashes'][suffix]:
                raise ValueError(f'Changed collected dataset file: {path}')
        sample_files.append(tuple(map(str, pair)))
    random.shuffle(sample_files)
    split = int(.8 * len(sample_files))
    if split == 0 or split == len(sample_files):
        raise ValueError('Need at least two collected instances for nonempty train/validation sets.')
    train_files, valid_files = sample_files[:split], sample_files[split:]
    if args.task == 'IP':
        from GCN import GNNPolicy_position as GNNPolicy, GraphDataset_position as GraphDataset
    else:
        from GCN import GNNPolicy, GraphDataset
    options = dict(batch_size=args.batch_size, num_workers=args.num_workers,
                   exclude_keys=['varInds', 'varNames'])
    train_loader = torch_geometric.loader.DataLoader(
        GraphDataset(train_files, max_solutions=args.max_solutions), shuffle=True, **options)
    valid_loader = torch_geometric.loader.DataLoader(
        GraphDataset(valid_files, max_solutions=args.max_solutions), shuffle=False, **options)
    model = GNNPolicy().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
    weight_norm = EnergyWeightNorm(args.task)
    if args.temperature is not None:
        weight_norm = (-1 if weight_norm < 0 else 1) * args.temperature
    output = root / 'pretrain' / f'{args.task}_train'
    log_dir = root / 'train_logs' / f'{args.task}_train'
    output.mkdir(parents=True, exist_ok=True)
    log_dir.mkdir(parents=True, exist_ok=True)
    if any(output.glob('*.pth')):
        raise ValueError('Checkpoints already exist; archive the previous run before retraining.')
    config = dict(vars(args), effective_device=str(device), energy_weight_norm=weight_norm,
                  torch_version=torch.__version__, pyg_version=torch_geometric.__version__,
                  train_instances=[Path(pair[0]).name[:-3] for pair in train_files],
                  valid_instances=[Path(pair[0]).name[:-3] for pair in valid_files], collection=collection)
    config_path = output / 'config.json'
    if config_path.exists() and json.loads(config_path.read_text()) != config:
        raise ValueError('Training configuration changed; use another --dataDir or archive the previous run.')
    # Engineering: provenance prevents graph/checkpoint mismatches; plain state_dict
    # files preserve compatibility with the original evaluation/checkpoint format.
    temporary_config = config_path.with_suffix('.json.tmp')
    temporary_config.write_text(json.dumps(config, indent=2, allow_nan=False) + '\n')
    temporary_config.replace(config_path)
    best_val_loss = float('inf')
    with (log_dir / f'{args.task}_train_train.log').open('w') as log_file:
        for epoch in range(args.epochs):
            begin = time.perf_counter()
            train_loss = train(model, train_loader, optimizer, weight_norm, device)
            valid_loss = train(model, valid_loader, None, weight_norm, device)
            if valid_loss < best_val_loss:
                best_val_loss = valid_loss
                temporary = output / 'model_best.pth.tmp'
                torch.save(model.state_dict(), temporary)
                temporary.replace(output / 'model_best.pth')
            temporary = output / 'model_last.pth.tmp'
            torch.save(model.state_dict(), temporary)
            temporary.replace(output / 'model_last.pth')
            line = (f'@epoch{epoch}   Train loss:{train_loss}   Valid loss:{valid_loss}'
                    f'    TIME:{time.perf_counter()-begin}\n')
            print(line, end='', flush=True)
            log_file.write(line)
            log_file.flush()
    print('done')


if __name__ == '__main__':
    main()