From 1c549eccb696f1e11e2a85db0d63b987adad41af Mon Sep 17 00:00:00 2001 From: naleraphael Date: Sat, 6 Jun 2020 22:25:33 +0800 Subject: [PATCH] TST: parameter `batch_size` was not passed into `DataLoader` This mistake made batch size of every data loader become the default value: 1. Though it does not affect the correctness of all test case, it still needs to be corrected. However, `batch_size` of a `DataLoader` cannot be modified after it is initialized. Therefore, we can only determine it while generating tasks for test, and that's why `batch_size` and `steps` is moved to the signature of `__init__` of each `Task`. --- tests/task.py | 62 ++++++++++++++++++++++++++++++++++----------------- 1 file changed, 42 insertions(+), 20 deletions(-) diff --git a/tests/task.py b/tests/task.py index 5e8d9b4..c79dab1 100644 --- a/tests/task.py +++ b/tests/task.py @@ -50,18 +50,25 @@ def __post_init__(self): class XORTask(BaseTask): - def __init__(self, validate=False): + def __init__(self, batch_size=8, steps=100, validate=False): super(XORTask, self).__init__() - bs, steps = 8, 64 - dataset = XORDataset(bs * steps) + n_total = batch_size * steps + dataset = XORDataset(n_total) if validate: - self.train_loader = DataLoader(Subset(dataset, range(steps - bs))) - self.val_loader = DataLoader(Subset(dataset, range(steps - bs, steps))) + n_train = int(n_total * 0.9) + self.train_loader = DataLoader( + Subset(dataset, range(n_train)), + batch_size=batch_size + ) + self.val_loader = DataLoader( + Subset(dataset, range(n_train, n_total)), + batch_size=batch_size + ) else: - self.train_loader = DataLoader(dataset) + self.train_loader = DataLoader(dataset, batch_size=batch_size) self.val_loader = None - self.batch_size = bs + self.batch_size = batch_size self.model = LinearMLP([8, 4, 1]) self.optimizer = optim.SGD(self.model.parameters(), lr=1e-5) self.criterion = nn.MSELoss() @@ -69,17 +76,25 @@ def __init__(self, validate=False): class ExtraXORTask(BaseTask): - def __init__(self, validate=False): + def __init__(self, batch_size=8, steps=100, validate=False): super(ExtraXORTask, self).__init__() - bs, steps = 8, 64 - dataset = ExtraXORDataset(bs * steps, extra_dims=2) + n_total = batch_size * steps + dataset = ExtraXORDataset(n_total, extra_dims=2) if validate: - self.train_loader = DataLoader(Subset(dataset, range(steps - bs))) - self.val_loader = DataLoader(Subset(dataset, range(steps - bs, steps))) + n_train = int(n_total * 0.9) + self.train_loader = DataLoader( + Subset(dataset, range(n_train)), + batch_size=batch_size + ) + self.val_loader = DataLoader( + Subset(dataset, range(n_train, n_total)), + batch_size=batch_size + ) else: - self.train_loader = DataLoader(dataset) + self.train_loader = DataLoader(dataset, batch_size=batch_size) self.val_loader = None + self.batch_size = batch_size self.model = LinearMLP([8, 4, 1]) self.optimizer = optim.SGD(self.model.parameters(), lr=1e-5) self.criterion = nn.MSELoss() @@ -87,18 +102,25 @@ def __init__(self, validate=False): class DiscriminativeLearningRateTask(BaseTask): - def __init__(self, validate=False): + def __init__(self, batch_size=8, steps=100, validate=False): super(DiscriminativeLearningRateTask, self).__init__() - bs, steps = 8, 64 - dataset = XORDataset(bs * steps) + n_total = batch_size * steps + dataset = XORDataset(n_total) if validate: - self.train_loader = DataLoader(Subset(dataset, range(steps - bs))) - self.val_loader = DataLoader(Subset(dataset, range(steps - bs, steps))) + n_train = int(n_total * 0.9) + self.train_loader = DataLoader( + Subset(dataset, range(n_train)), + batch_size=batch_size + ) + self.val_loader = DataLoader( + Subset(dataset, range(n_train, n_total)), + batch_size=batch_size + ) else: - self.train_loader = DataLoader(dataset) + self.train_loader = DataLoader(dataset, batch_size=batch_size) self.val_loader = None - dataset = XORDataset(128) + self.batch_size = batch_size self.model = LinearMLP([8, 4, 1]) self.optimizer = optim.SGD( [