| 85 | |
| 86 | |
| 87 | class CreateDataset(Dataset): |
| 88 | def __init__(self, occ, prc, lb, pt, device, adj): # adj |
| 89 | occ, label = create_rnn_data(occ, lb, pt) |
| 90 | prc, _ = create_rnn_data(prc, lb, pt) |
| 91 | self.occ = torch.Tensor(occ) |
| 92 | self.prc = torch.Tensor(prc) |
| 93 | self.label = torch.Tensor(label) |
| 94 | self.device = device |
| 95 | |
| 96 | def __len__(self): |
| 97 | return len(self.occ) |
| 98 | |
| 99 | def __getitem__(self, idx): # occ: batch, seq, node |
| 100 | output_occ = torch.transpose(self.occ[idx, :, :], 0, 1).to(self.device) |
| 101 | output_prc = torch.transpose(self.prc[idx, :, :], 0, 1).to(self.device) |
| 102 | output_label = self.label[idx, :].to(self.device) |
| 103 | return output_occ, output_prc, output_label |
| 104 | |
| 105 | |
| 106 | class CreateFastDataset(Dataset): |
nothing calls this directly
no outgoing calls
no test coverage detected