calibrating.py

# Borrowed from NAF. Might be plugged into DeepSets.

class SigmoidFlow(nn.Module):
    def __init__(self,K):
        super().__init__()
        self.W = nn.Parameter(torch.FloatTensor(K,1))
        self.b = nn.Parameter(torch.FloatTensor(1,K))
        self.alpha = nn.Linear(1,K)
        
    def forward(self,x):
        alpha = F.softmax(self.alpha(x),dim=-1)
        x = F.linear(x,F.softplus(self.W),self.b)
        x = F.sigmoid(x)
        x = (alpha * x).sum(dim=-1).unsqueeze(-1)
        return x
    
class Calibrator(nn.Module):
    def __init__(self,layers=2,K=5):
        super().__init__()
        self.inner = nn.Sequential(
            *[
                SigmoidFlow(K=K) for _ in range(layers)
            ]
        )
        
    def forward(self,x):
        return self.inner(x)