Pyro model/guide neural net usage

pumplerod: <p>I had originally created <a href="https://stackoverflow.com/questions/76196846/pyro-model-guide-neural-net-usage" rel="noopener nofollow ugc">this post</a> on <code>stackoverflow</code>, though I should probably posed the question here first.</p>
<p>I’ve been struggling with understanding/implementing a pyTorch neural net model using pyro.  Examples from the documentation don’t exactly paint a picture for what is going on inside the <code>Predictive</code> model.  I’ve been able to use the examples to get a functioning model/guide however when I try to implement for my use case it’s not clear how to adjust in order to properly implement a solution.</p>
<p>Basic sticking points…</p>
<ol>
<li>How does the <code>Predictive</code> model use both the <code>model</code> and <code>guide</code>?</li>
<li>Examples using <code>Predictive</code> show retrieving the result from the <code>obs</code> sample (which is in the model, not the guide).</li>
<li>Since the <code>obs</code> sample is in fact a <code>pyro.sample</code> object, should I simply draw N samples from the <code>obs</code> sample?  Or should I run the model N times to generate all the results since the weights/bias values for my network are also <code>pyro.sample</code> objects?</li>
</ol>
<p>As an example…</p>
<p>I have a dataset with features and targets grouped into discrete bins <code>[0.0, 0.25, 0.5, 0.75, 1.0]</code></p>
<p>I would like to categorize the samples to accurately determine if a sample’s target value is either <code>0.0</code> or <code>1.0</code> (ignoring the other target values for the time being).  Ultimately I want to use the logits from the model prediction to determine if a sample is a <code>0.0</code> or a <code>1.0</code> or use the logits as a confidence value and determine a single continuous value between the two targets.</p>
<p>When no longer training and performing evaluation I would like to normalize the features for the batch of samples, so I have a normalization performed if the model is not in <code>train</code> mode.</p>
<p>I would then like my final prediction to be a continuous value between <code>0.0-1.0</code>.  I am really only interested in accurately predicting the 0.0 and 1.0 values, however I would then like to establish the other values as a likelihood between 0.0-1.0, so I am using a OneHot encoding with 2 values per output.  [1.0, 0.0] being the likelihood my sample is a 0.0 and the [0.0, 1.0] being a likelihood my sample is a 1.0.  I then average all the votes from <code>obs</code> to determine where this sample lies between 0 and 1.</p>
<p>While this seems to work, I feel like I’m probably being very inefficient.</p>
<p>After training, why do I need to create a <code>Predictive</code> instance of my model?  since the weights and bias of the network are also sampled parameters, could I also simply run a prediction multiple times?  In the <code>Predictive</code> instance, what is actually being used for the weight/bias values of the network.  Do the come from the <code>model</code> or do they come from the <code>guide</code>?</p>
<p>Here I am trying to make use of <code>pyro.nn.PyroModule</code>, however I had trouble getting this to work with <code>AutoGuide</code> so I’m attempting to create my own <code>guide()</code> method within my model.</p>
<pre><code class="lang-auto">import torch
import pyro
from pyro.infer import SVI, Trace_ELBO, Predictive

# Pyro Linear Layer Class 
#   Used to create a fully connected linear layer with Normally distributed weights
#   and bias where mu/sigma become a learnable parameter
class PyroLinear(torch.nn.Linear, pyro.nn.PyroModule):  # used as a mixin
    def __init__(self, in_features, out_features, device='cpu', **kwargs):
        super().__init__(in_features, out_features, **kwargs)

        mu = torch.randn_like( self.weight, device=device)
        sigma = torch.rand_like( self.weight, device=device) + 0.01
        self.weight = pyro.nn.PyroSample( pyro.distributions.Normal( mu, sigma).expand([self.out_features, self.in_features]).to_event(2))

        if self.bias is not None:
            mu = torch.randn_like( self.bias, device=device)
            sigma = torch.rand_like( self.bias, device=device) + 0.01
            self.bias = pyro.nn.PyroSample( pyro.distributions.Normal( mu, sigma).expand([self.out_features]).to_event(1))

#
# helper function to return a fully connected pyro layer with ReLU activation
def layer_block( n_in, n_out, device='cpu'):
    return pyro.nn.PyroModule[ torch.nn.Sequential](
        PyroLinear(n_in, n_out, device=device),
        torch.nn.ReLU()
    )

#
# Full NNet model with guide method
class NNet_Model(pyro.nn.PyroModule):
    def __init__(self, n_inputs=1, n_outputs=1, h_layers=[20], device='cpu'):
        super().__init__()
        self.n_outputs = n_outputs
        self.layer_sizes = [n_inputs, *h_layers]

        layer_blocks = [ layer_block(in_f, out_f, device=device) for in_f, out_f in zip(self.layer_sizes[:-1], self.layer_sizes[1:])]
        
        self.feature_net = pyro.nn.PyroModule[ torch.nn.Sequential](*layer_blocks)
        self.out = PyroLinear( self.layer_sizes[ -1], n_outputs, device=device)

    def forward(self, x, y=None):
        #
        # Normalize the batch if this is being run for an eval() prediction
        if not self.training:
            x = (x - x.mean( axis=0)) / (x.std( axis=0) + 1e-6)
            
        x = self.feature_net( x)
        pred_logits = self.out( x)

        obs = pyro.sample("obs", pyro.distributions.OneHotCategorical(logits=pred_logits).to_event(1), obs=y)
        return pred_logits
        
    #
    # The guide.  This seems to work even though it does not return anything
    #    - I'm not sure what it should return since the model is generating 
    #      the pred_logits.
    def guide(self, x, y=None):
        #
        # Create a pyro.param for each mu/sigma of the weights/bias of feature_net
        for i in range(0, len( self.feature_net)):
            mu = torch.randn_like( self.feature_net[i][0].weight)
            sigma = torch.rand_like( self.feature_net[i][0].weight) + 0.1
            mu_param = pyro.param( f"feature_net.{i}.0.w_mu", mu)
            sigma_param = pyro.param( f"feature_net.{i}.0.w_sigma", sigma, constraint=pyro.distributions.constraints.positive)
            _ = pyro.sample( f"feature_net.{i}.0.weight",  pyro.distributions.Normal(mu_param, sigma_param).expand([self.feature_net[i][0].weight.size(0), self.feature_net[i][0].weight.size(1)]).to_event(2))

            mu = torch.randn_like( self.feature_net[i][0].bias)
            sigma = torch.rand_like( self.feature_net[i][0].bias) + 0.1
            mu_param = pyro.param( f"feature_net.{i}.0.b_mu", mu)
            sigma_param = pyro.param( f"feature_net.{i}.0.b_sigma", sigma, constraint=pyro.distributions.constraints.positive)
            _ = pyro.sample( f"feature_net.{i}.0.bias",  pyro.distributions.Normal(mu_param, sigma_param).expand([self.feature_net[i][0].bias.size(0)]).to_event(1))
        #
        # Create pyro.param for the mu/sigma of the weights/bias of the output layer
        mu = torch.randn_like( self.out.weight)
        sigma = torch.rand_like( self.out.weight) + 0.1
        mu_param = pyro.param( f"out.w_mu", mu)
        sigma_param = pyro.param( f"out.w_sigma", sigma, constraint=pyro.distributions.constraints.positive)
        _ = pyro.sample( f"out.weight",  pyro.distributions.Normal(mu_param, sigma_param).expand([self.out.weight.size(0), self.out.weight.size(1)]).to_event(2))

        mu = torch.randn_like( self.out.bias)
        sigma = torch.rand_like( self.out.bias) + 0.1
        mu_param = pyro.param( f"out.b_mu", mu)
        sigma_param = pyro.param( f"out.b_sigma", sigma, constraint=pyro.distributions.constraints.positive)
        _ = pyro.sample( f"out.bias",  pyro.distributions.Normal(mu_param, sigma_param).expand([self.out.bias.size(0)]).to_event(1))
</code></pre>
<p>Here I create a simple set of fake training data…</p>
<pre data-code-wrap="classes"><code class="lang-classes">n_features = 9
n_samples = 1000
DEVICE='cpu'

feature_data = torch.tensor( np.random.choice( classes, (n_samples, n_features))).type( torch.float32)#.to( DEVICE)
f_idxs = np.arange(0, n_features)
t_build = feature_data[ :, np.random.choice( f_idxs, (1))] - feature_data[ :, np.random.choice( f_idxs, (1))] + feature_data[ :, np.random.choice( f_idxs, (1))]
target_data = torch.maximum( torch.minimum( t_build, torch.tensor( 1.0)), torch.tensor( 0.0))#.to( DEVICE)

#
# Split into train/test groups.
#  - normalize the training feature data
train_feature_data = feature_data[ 0:800, :]
train_feature_data = (train_feature_data - train_feature_data.mean( axis=0)) / ( train_feature_data.mean( axis=0) + 1e-6)
train_target_data = target_data[ 0:800, :]
test_feature_data = feature_data[ 800:, :]
test_target_data = target_data[ 800:, :]

#
# Now select only the data with target values of 0.0 or 1.0
train_feature_01_data[ ((train_target_data==1.0) | (train_target_data==0.0)).view( -1), :]
train_target_01_data[ ((train_target_data==1.0) | (train_target_data==0.0)).view( -1), :]
</code></pre>
<p>Train model on the <code>train_feature_01_data</code> and <code>train_target_01_data</code>…</p>
<pre><code class="lang-auto">pyro.clear_param_store()
#
# I am using OneHot encoding, however I am only interested in predicting the 0,1 values so I am only outputting 2 values.
model = NNet_Model( n_inputs=n_features, n_outputs=2, h_layers=[ 10000], device=DEVICE)

adam = pyro.optim.Adam({"lr": 4e-3})
svi = SVI(model, guide=model.guide, optim=adam, loss=Trace_ELBO())

y = torch.nn.functional.one_hot( (train_target_data.view(-1)).type(torch.long), num_classes=2)

losses = []
for epoch in range(0, 200000):
    loss = svi.step( train_feature_01_data, y)
    losses.append( loss)
    print(f'loss={loss:.4f}', end='\r')
</code></pre>
<p>After training, I create a <code>Predictive</code> instance of my model to test on the <code>test_feature_data</code> against the <code>test_target_data</code> which comprise all target values [0.0, 0.25, 0.5, 0.75, 1.0]</p>
<pre><code class="lang-auto">predictive_model = Predictive(model, guide=model.guide, num_samples=1000)
preds = predictive_model( test_feature_data)
</code></pre>
<p>now <code>preds[ 'obs']</code> contains 1000 OneHot predictions for each sample.  To boil it down to a final prediction I average the predictions per sample ( using <code>argmax</code> to determine whether the sample is a 0.0 label or a 1.0 label…</p>
<pre><code class="lang-auto">final_pred = preds[ 'obs'].argmax( axis=2).type( torch.float).mean(axis=0)
</code></pre>
<p>I then sort the values based on the <code>test_target_data</code> so that it is easier to identify on a line chart.</p>
<pre><code class="lang-auto">values, idxs = test_target_data.view( -1).sort()
plt.plot( test_target_data.view( -1)[ idxs].numpy(), label="target")
plt.plot( final_pred.view( -1)[ idxs].numpy(), label="Prediction")
plt.legend()
</code></pre>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/5/9576d603468d9289b8f81f619ec5774ad15f15db.png" alt="pyro_example" data-base62-sha1="lkdKnpvugMVSKhjUw8tlO3ual1p" width="374" height="249"></p>
<p>As you can see, for the most part it does seem to organize even the middle values to their proper space, though it has trouble with targets that are in the 0.75 range.  But this is just my random simple test data.</p>
<p>I would very much appreciate any clues as to how to properly build/train/save a <code>pyro.nn.PyroModule</code> style model.  I have found it tricky to get it to/from the GPU and, ultimately, for the final <code>Predictive</code> model I need to run on the CPU because my GPU cannot handle the <code>obs</code> tensor with 1000 samples per row of data.  My workarounds for this seem to do the trick, but it feels clumsy.</p>
<p>Mainly I would like to know what the best method for saving and re-using a model is once it has gone through training.</p>
<p>is generating a <code>Predictive</code> model from the <code>original model</code> and <code>guide</code> essentially the same as generating predictions n-times from the <code>original model</code>?  How does the <code>guide</code> factor into the final predicting process?</p>