SymILO/dataset.py at main · NetSysOpt/SymILO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import numpy as np
import copy
from feature_extractor import extractFeature
import torch
from torch.utils.data import Dataset
import os
import gzip
import pickle
from config import *

import argparse


class MIPDataset(Dataset):
    def __init__(self,files,bgdir,reorderFunc, addPosFuc):
        insPaths = [ filepaths[0] for filepaths in files]
        solPaths = [ filepaths[1] for filepaths in files]
        self.insPaths = insPaths
        self.solPaths = solPaths
        self.bgdir = bgdir
        self.reorder = reorderFunc
        self.addPos = addPosFuc
        os.makedirs(bgdir,exist_ok=True)

    def __getitem__(self, index):
        inspath = self.insPaths[index]
        solpath = self.solPaths[index]

        insname = os.path.basename(inspath)
        bgpath = os.path.join(self.bgdir,insname+'.bg')
        if os.path.exists(bgpath):
            data = pickle.load(gzip.open(bgpath,'rb'))
        else:
            features = extractFeature(inspath)
            features = self.addPos(features)
            varNames = np.array(features.varNames)[features.biInds]
            reorderData = self.reorder(varNames)
            data = {
                'groupFeatures':torch.Tensor(features.groupFeatures),
                'varFeatures': torch.Tensor(features.varFeatures),
                'consFeatures':torch.Tensor(features.consFeatures),
                'edgeFeatures':torch.Tensor(features.edgeFeatures),
                'edgeInds':torch.Tensor(features.edgeInds.astype(int)).permute(1,0),
                'biInds':torch.Tensor(features.biInds).long(),
                'nGroup':reorderData['nGroup'],
                'nElement':reorderData['nElement'],
                'reorderInds':torch.Tensor(reorderData['reorderInds'])
            }

            if self.solPaths[index] is not None:
                solData = pickle.load(gzip.open(solpath, 'rb'))
                sols = solData['sols']
                objs = solData['objs']
                varNames = solData['varNames']


                varIds = list(range(len(varNames)))
                varTuples = list(zip(varNames, varIds))
                varTuples.sort(key=lambda t: t[0])
                order = [t[-1] for t in varTuples]

                sols = sols[:,order][:,features.biInds]


                data['sols'] = torch.Tensor(sols[0])
                data['objs'] = torch.Tensor([objs[0]])

            pickle.dump(data,gzip.open(bgpath,'wb'))


        return data

    def __len__(self):
        return len(self.insPaths)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--dataset', type=str, default='SMSP')
    args = parser.parse_args()
    info = confInfo[args.dataset]
    ADDPOS = info['addPosFeature']
    REORDER = info['reorder']

    for mod in ['train','test']:

        fileDir = os.path.join(info[f'{mod}Dir'], 'ins')
        solDir = os.path.join(info[f'{mod}Dir'], 'sol')
        bgDir = os.path.join(info[f'{mod}Dir'], 'bg')

        filenames = os.listdir(fileDir)
        filepaths = [os.path.join(fileDir, filename) for filename in filenames]
        solpaths = [os.path.join(solDir, filename+'.sol') if mod == 'train' else None for filename in filenames ]
        dataset = MIPDataset(list(zip(filepaths,solpaths)),bgDir,REORDER,ADDPOS)

        data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=0)
        print(f'Start constructing bipartite graph for {mod} set ...')
        for step,data in enumerate(data_loader):
            print(f'Processed {step}/{len(data_loader)}')
        print(f'Bipartite graph construction for {mod} set finished!')