Masking and Instance Normalization in PyTorch
Assume I have a PyTorch tensor, arranged as shape [N, C, L] where N is the batch size, C is the number of channels or features, and L is the length.  In this case, if one wishes to perform instance normalization, one does something like:

N = 20
C = 100
L = 40
m = nn.InstanceNorm1d(C, affine=True)
input = torch.randn(N, C, L)
output = m(input) 


This will perform a normalization in the L-wise dimension for each N*C = 2000 slices of data, subtracting 2000 means, scaling by 2000 standard deviations, and re-scaling by 100 learnable weight and bias parameters (one per channel).  The unspoken assumption here is that all of these values exist and are meaningful.  

But I have a situation where, for the slice N=1, I would like to exclude all data after (say) L=35.  For the slice N=2 (say) all the data are valid.  For the slice N=3, exclude all data after L=30, etc.  This mimics data which are one dimensional time sequences, having multiple features, but which are not the same length.

How can I perform an instance norm on such data, get correct statistics, and maintain differentiability/AutoGrad information in PyTorch?

Update:  While maintaining GPU performance, or at least not killing it dead.

I cannot...


...Mask with zero values, as this destroys the computer means and variances giving erroneous results
...Mask with np.nan or np.inf, as PyTorch tensors do not ignore such values, but treat them as errors.  They are sticky, and lead to garbage results.  PyTorch currently lacks the equivalent of np.nanmean and np.nanvar.
...Permute or transpose to an amenable arrangement of data; no such approach gives me what I need
...Use a pack_padded_sequence; instance normalization does not operate on that data structure, and one cannot import data into that structure as far as I know.  Also, data re-arrangement would still be necessary, see 3 above. 


Am I missing an approach which would give me what I need?  Or perhaps am I missing a method of data re-arrangement which would allow 3 or 4 above to work?

This is an issue faced by recurrent neural networks all the time, hence the pack_padded_sequence functionality, but it isn't quite applicable here. 

I don't think this is directly possible to implement using the existing InstanceNorm1d, the easiest way would probably be implementing it yourself from scratch. I did a quick implementation that should work. To make it a little bit more general this module requires a boolean mask (a boolean tensor of the same size as the input) that specifies which elements should be considered when passing through the instance norm.

import torch


class MaskedInstanceNorm1d(torch.nn.Module):
    def __init__(self, num_features, eps=1e-6, momentum=0.1, affine=True, track_running_stats=False):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        self.affine = affine
        self.track_running_stats = track_running_stats
        self.gamma = None
        self.beta = None
        if self.affine:
            self.gamma = torch.nn.Parameter(torch.ones((1, self.num_features, 1), requires_grad=True))
            self.beta = torch.nn.Parameter(torch.zeros((1, self.num_features, 1), requires_grad=True))

        self.running_mean = None
        self.running_variance = None
        if self.affine:
            self.running_mean = torch.zeros((1, self.num_features, 1), requires_grad=True)
            self.running_variance = torch.zeros((1, self.num_features, 1), requires_grad=True)

    def forward(self, x, mask):
        mean = torch.zeros((1, self.num_features, 1), requires_grad=False)
        variance = torch.ones((1, self.num_features, 1), requires_grad=False)

        # compute masked mean and variance of batch
        for c in range(self.num_features):
            if mask[:, c, :].any():
                mean[0, c, 0] = x[:, c, :][mask[:, c, :]].mean()
                variance[0, c, 0] = (x[:, c, :][mask[:, c, :]] - mean[0, c, 0]).pow(2).mean()

        # update running mean and variance
        if self.training and self.track_running_stats:
            for c in range(self.num_features):
                if mask[:, c, :].any():
                    self.running_mean[0, c, 0] = (1-self.momentum) * self.running_mean[0, c, 0] \
                                                 + self.momentum * mean[0, c, 0]
                    self.running_variance[0, c, 0] = (1-self.momentum) * self.running_variance[0, c, 0] \
                                                     + self.momentum * variance[0, c, 0]

        # compute output
        x = (x - mean)/(self.eps + variance).sqrt()

        if self.affine:
            x = x * self.gamma + self.beta

        return x

