What is the state of the art way of doing regression with probability in pytorch
All regression examples I find are examples where you predict a real number and unlike with classification you dont the the confidence the model had when predicting that number. I have done in reinforcement learning another way the output is instead the mean and std and then you sample from that distribution. Then you know how confident the model is at predicting every value. Now I cant find how to do this using supervised learning in pytorch. The problem is that I dont understand how to perform sample from the distribution the get the actual value while training or what sort of loss function I should use, not sure how for example MSE or L1Smooth would work.

Is there any example ot there where this is done in pytorch in a robust and state of the art way?

The key point is that you do not need to sample from the NN-produced distribution. All you need is to optimize the likelihood of the target value under the NN distribution.

There is an example in the official PyTorch example on VAE (https://github.com/pytorch/examples/tree/master/vae), though for multidimensional Bernoulli distribution.

Since PyTorch 0.4, you can use torch.distributions: instantiate distribution  distro with outputs of your NN and then optimize  -distro.log_prob(target).

EDIT: As requested in a comment, a complete example of using the torch.distributions module.

First, we create a heteroscedastic dataset:

import numpy as np
import torch
X = np.random.uniform(size=300)
Y = X + 0.25*X*np.random.normal(size=X.shape[0])


We build a trivial model, which is perfectly able to match the generative process of our data:

class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.mean_coeff = torch.nn.Parameter(torch.Tensor([0]))
        self.var_coeff = torch.nn.Parameter(torch.Tensor([1]))

    def forward(self, x):
        return torch.distributions.Normal(self.mean_coeff * x, self.var_coeff * x)

mdl = Model()
optim = torch.optim.SGD(mdl.parameters(), lr=1e-3)


Initialization of the model makes it always produce a standard normal, which is a poor fit for our data, so we train (note it is a very stupid batch training, but demonstrates that you can output a set of distributions for your batch at once):

for _ in range(2000): # epochs
    dist = mdl(torch.from_numpy(X).float())
    obj = -dist.log_prob(torch.from_numpy(Y).float()).mean()
    optim.zero_grad()
    obj.backward()
    optim.step()


Eventually, the learned parameters should match the values we used to construct the Y.

print(mdl.mean_coeff, mdl.var_coeff)
# tensor(1.0150) tensor(0.2597)

