Get dataloaders for training and testing Parameters ---------- batch_size: int Batch size for training test_batch_size: int Batch size for testing Returns ------- dict Dictionary of dataloaders
(self, batch_size, test_batch_size = None)
| 1169 | print_sys("Done!") |
| 1170 | |
| 1171 | def get_dataloader(self, batch_size, test_batch_size = None): |
| 1172 | """ |
| 1173 | Get dataloaders for training and testing |
| 1174 | |
| 1175 | Parameters |
| 1176 | ---------- |
| 1177 | batch_size: int |
| 1178 | Batch size for training |
| 1179 | test_batch_size: int |
| 1180 | Batch size for testing |
| 1181 | |
| 1182 | Returns |
| 1183 | ------- |
| 1184 | dict |
| 1185 | Dictionary of dataloaders |
| 1186 | |
| 1187 | """ |
| 1188 | if test_batch_size is None: |
| 1189 | test_batch_size = batch_size |
| 1190 | |
| 1191 | self.node_map = {x: it for it, x in enumerate(self.adata.var.gene_name)} |
| 1192 | self.gene_names = self.adata.var.gene_name |
| 1193 | |
| 1194 | # Create cell graphs |
| 1195 | cell_graphs = {} |
| 1196 | if self.split == 'no_split': |
| 1197 | i = 'test' |
| 1198 | cell_graphs[i] = [] |
| 1199 | for p in self.set2conditions[i]: |
| 1200 | if p != 'ctrl': |
| 1201 | cell_graphs[i].extend(self.dataset_processed[p]) |
| 1202 | |
| 1203 | print_sys("Creating dataloaders....") |
| 1204 | # Set up dataloaders |
| 1205 | test_loader = DataLoader(cell_graphs['test'], |
| 1206 | batch_size=batch_size, shuffle=False) |
| 1207 | |
| 1208 | print_sys("Dataloaders created...") |
| 1209 | return {'test_loader': test_loader} |
| 1210 | else: |
| 1211 | if self.split =='no_test': |
| 1212 | splits = ['train','val'] |
| 1213 | else: |
| 1214 | splits = ['train','val','test'] |
| 1215 | for i in splits: |
| 1216 | cell_graphs[i] = [] |
| 1217 | for p in self.set2conditions[i]: |
| 1218 | cell_graphs[i].extend(self.dataset_processed[p]) |
| 1219 | |
| 1220 | print_sys("Creating dataloaders....") |
| 1221 | |
| 1222 | # Set up dataloaders |
| 1223 | train_loader = DataLoader(cell_graphs['train'], |
| 1224 | batch_size=batch_size, shuffle=True, drop_last = True) |
| 1225 | val_loader = DataLoader(cell_graphs['val'], |
| 1226 | batch_size=batch_size, shuffle=True) |
| 1227 | |
| 1228 | if self.split !='no_test': |