ig248 / deepquantiles Goto Github PK

def fit(self, data: numpy.ndarray, targets: numpy.ndarray = None) -> 'FBUncertaintyRegressor':
    data = torch.tensor(data, dtype=torch.float32)
    targets = torch.tensor(targets, dtype=torch.float32).to(device=self.device)
    ds = Data.TensorDataset(data, targets)
    loader = Data.DataLoader(dataset=ds, batch_size=self.batch_size)

    dim = data.shape[1]
    self.model = torch.nn.Sequential(
        torch.nn.Linear(dim + 1, self.hidden),
        torch.nn.ReLU(),
        torch.nn.BatchNorm1d(self.hidden),
        torch.nn.Dropout(.2),
        torch.nn.Linear(self.hidden, 1)
    ).to(device=self.device)

    opt = torch.optim.Adam(self.model.parameters(), eps=1e-07,
                           lr=self.learning_rate)   #, weight_decay=self.weight_decay)
    loss = QuantileLoss()

    for i in range(self.epochs):
        for batch_x, batch_y in loader:
            batch_x = batch_x.to(device=self.device)
            opt.zero_grad()
            if self.quantil == "all":
                taus = torch.rand(batch_x.shape[0], 1).to(device=self.device)
            else:
                taus = torch.zeros(batch_x.shape[0], 1).fill_(self.quantil).to(device=self.device)
            tau_augs = self.__augment__(batch_x, taus)
            model_out = self.model(tau_augs)
            a = loss(model_out, targets, taus)
            a.backward()
            opt.step()
        print(i, a)
        #taus_ = torch.rand(data.size(0), 1).fill_(self.quantil).to(device=self.device)
        #loss_ = loss(self.model(self.__augment__(data, taus_)), targets, taus_).detach().numpy()
        #print(f"Epoch {i}/{self.epochs}; loss: {loss_}")

    return self

def predict(self, data: numpy.ndarray, predict_quantil: float = None) -> numpy.ndarray:
    if predict_quantil is None and isinstance(self.quantil, float):
        predict_quantil = self.quantil
    data = torch.tensor(data, dtype=torch.float32).to(device=self.device)
    return self.model(self.__augment__(data, predict_quantil)).to(device='cpu').detach().numpy()

def __augment__(self, data: torch.Tensor, tau=None) -> torch.Tensor:
    if tau is None:
        tau = torch.zeros(data.size(0), 1).fill_(0.5).to(device=self.device)
    elif type(tau) == float:
        tau = torch.zeros(data.size(0), 1).fill_(tau).to(device=self.device)
    return torch.cat((data, (tau - 0.5) * 12), 1)

def forward(self, yhat: torch.Tensor, y: torch.Tensor, tau: float) -> torch.Tensor:
    diff = yhat - y
    mask = (diff.ge(0).float() - tau).detach()
    return (mask * diff).mean()