MCPcopy Create free account
hub / github.com/Modulus-Labs/RockyBot / Playground_Dataset

Class Playground_Dataset

pytorch-model/datasets.py:280–331  ·  view source on GitHub ↗

Subject to change!

Source from the content-addressed store, hash-verified

278
279
280class Playground_Dataset(Dataset):
281 """
282 Subject to change!
283 """
284 def __init__(self, mode):
285 if mode == "train":
286 print(f"Setting up train playground dataset...")
287 self.data_path = constants.get_dataset_filepath(
288 constants.PLAYGROUND_TASK, constants.TRAIN_DATASET_FILENAME)
289 self.label_path = constants.get_dataset_filepath(
290 constants.PLAYGROUND_TASK, constants.TRAIN_LABELS_FILENAME)
291 elif mode == "val":
292 print(f"Setting up val playground dataset...")
293 self.data_path = constants.get_dataset_filepath(
294 constants.PLAYGROUND_TASK, constants.VAL_DATASET_FILENAME)
295 self.label_path = constants.get_dataset_filepath(
296 constants.PLAYGROUND_TASK, constants.VAL_LABELS_FILENAME)
297 else:
298 print(f"Error: mode should be one of [train, val] but got {mode} instead.")
299
300 # --- Load in dataset + labels ---
301 with open(self.data_path, "rb") as f:
302 self.x = torch.from_numpy(np.load(f)).float()
303 with open(self.label_path, "rb") as f:
304 self.y = torch.from_numpy(np.load(f)).long()
305
306 self.num_classes = int(torch.max(self.y).item() + 1)
307
308 # --- Load in label semantics ---
309 self.features_to_idx_path = constants.get_dataset_filepath(
310 constants.PLAYGROUND_TASK, constants.FEATURES_TO_IDX_FILENAME)
311 self.labels_to_idx_path = constants.get_dataset_filepath(
312 constants.PLAYGROUND_TASK, constants.LABELS_TO_IDX_FILENAME)
313 with open(self.features_to_idx_path, "r") as f:
314 self.features_to_idx = json.load(f)
315 with open(self.labels_to_idx_path, "r") as f:
316 self.labels_to_idx = json.load(f)
317
318 def get_weights(self):
319 return get_weights(self.y)
320
321 def get_input_dim(self):
322 return self.x.shape[-1]
323
324 def get_output_dim(self):
325 return self.num_classes
326
327 def __getitem__(self, idx):
328 return self.x[idx], self.y[idx]
329
330 def __len__(self):
331 return len(self.x)
332
333
334# --- For argparse ---

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected