File size: 1,120 Bytes
52933b5 |
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 |
import logging
import pickle
root_dir = "/home/cwan5/OR/attention-learn-to-route/data/tsp/"
def load(filename, root_dir=root_dir):
return pickle.load(open(root_dir + filename, "rb"))
file_catalog = {
"test": {
20: "tsp20_test_seed1234.pkl",
50: "tsp50_test_seed1234.pkl",
100: "tsp100_test_seed1234.pkl",
},
"eval": {
20: "tsp20_validation_seed4321.pkl",
50: "tsp50_validation_seed4321.pkl",
100: "tsp100_validation_seed4321.pkl",
},
}
class lazyClass:
data = {
"test": {},
"eval": {},
}
def __getitem__(self, index):
partition, nodes, idx = index
if not (partition in self.data) or not (nodes in self.data[partition]):
logging.warning(
f"Data sepecified by ({partition}, {nodes}) was not initialized. Attepmting to load it for the first time."
)
data = load(file_catalog[partition][nodes])
self.data[partition][nodes] = [tuple(instance) for instance in data]
return self.data[partition][nodes][idx]
TSPDataset = lazyClass()
|