|
发表于 2021-11-10 21:38:53
|
显示全部楼层
dataset = DHCPDataset(file_path, regression=regression)
item_count = len(dataset)
train_num = int(item_count * ratio)
train_sampler = SubsetRandomSampler(list(range(0, train_num)))
test_sampler = SubsetRandomSampler(list(range(train_num, item_count)))
# init dataloader
train_loader = DataLoader(
dataset, batch_size=1, sampler=train_sampler, collate_fn=lambda x:x[0])
test_loader = DataLoader(
dataset, batch_size=1, sampler=test_sampler, collate_fn=lambda x:x[0]) |
|