class ExtraXORTask(BaseTask):
def __init__(self, validate=False):
super(ExtraXORTask, self).__init__()
bs, steps = 8, 64
dataset = ExtraXORDataset(bs * steps, extra_dims=2)
if validate:
self.train_loader = DataLoader(Subset(dataset, range(steps - bs)))
self.val_loader = DataLoader(
Subset(dataset, range(steps - bs, steps)))
else:
self.train_loader = DataLoader(dataset)
self.val_loader = None
self.model = LinearMLP([8, 4, 1])
self.optimizer = optim.SGD(self.model.parameters(), lr=1e-5)
self.criterion = nn.MSELoss()
self.device = torch.device("cuda")
class ExtraXORTask(BaseTask):
def __init__(self, batch_size=8, steps=100, validate=False):
super(ExtraXORTask, self).__init__()
n_total = batch_size * steps
dataset = ExtraXORDataset(n_total, extra_dims=2)
if validate:
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, 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()
self.device = torch.device("cuda")
