Neural network architecture suggestion(s) for video prediction (image sequence)

marco_c: <p>Hello everybody! Could you please provide me neural network architecture suggestion(s) for video prediction (image sequence) regarding entering 144 images and predicting 48 images for each sequence? For illustration, I tested the Wavenet and CNNLSTM networks contained in <a href="https://bitbucket.org/retiarus/tec_prediction/src/conv-lstm/model/networks.py" class="inline-onebox" rel="noopener nofollow ugc">Bitbucket</a> (but not only them) and I didn’t get good results. Thanks in advance!</p>

Kapil_Rana: <p>If you have good computation resources, You can try <strong>vision transformers</strong>. There are multiple papers in the literature.</p>

J_Johnson: <p>How I’d approach this is basically modify a 2d UNet to use all 3d layers, and possibly adding in a 4d branch internally to the model, to assist with modeling 3d objects/scenes as they change over time(3d+1d=4d).</p>
<p>Here is an example of a standard 2d Unet as well as a 2d + 3d Unet hybrid:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://github.com/CerebralSeed/Hybrid-3D-UNet">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" width="32" height="32">

      <a href="https://github.com/CerebralSeed/Hybrid-3D-UNet" target="_blank" rel="noopener nofollow ugc">GitHub</a>
  </header>

  <article class="onebox-body">
    <div class="aspect-image" style="--aspect-ratio:690/345;"><img src="https://opengraph.githubassets.com/50c6e1f73a90b1bd45c337894e702bfcc67ceb2e2b0bdba33685617f4819bc6f/CerebralSeed/Hybrid-3D-UNet" class="thumbnail" width="690" height="345"></div>

<h3><a href="https://github.com/CerebralSeed/Hybrid-3D-UNet" target="_blank" rel="noopener nofollow ugc">GitHub - CerebralSeed/Hybrid-3D-UNet: Model for Hybrid 3D UNet</a></h3>

  <p>Model for Hybrid 3D UNet. Contribute to CerebralSeed/Hybrid-3D-UNet development by creating an account on GitHub.</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p>You’ll want to ensure your “channels” are still RBG and not the time sequence dimension. Also, adjust the kernel sizes throughout so it accepts 144 into the encoder and gives 48 out on the decoder on the time sequence dim.</p>

marco_c: <p>Hi. Ok, I’ll look into that, thanks!</p>

marco_c: <p>Hi. Ok, I took a look into that…</p>
<p>I’m also currently testing a UNET-LSTM network according to the code below, but it also didn’t show good results so far, it would be very different from the UNET you suggested, besides the fact that it doesn’t have LSTM, right?</p>
<pre><code class="lang-auto"> def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_channels, out_channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

class UnetLSTM(nn.Module):
    def __init__(self, in_channels, out_channels, num_layers, kernel_size, dilation, stride, dropout):
        super(UnetLSTM, self).__init__()

        self.in_channels = in_channels  # Corresponds to input size
        self.out_channels = out_channels  # Corresponds to hidden size
        self.num_layers = num_layers 
        self.kernel_size = kernel_size
        self.dilation = dilation
        self.stride = stride
        self.dropout = dropout

        #self.conv1 = nn.Conv3d(in_channels=1, out_channels=128, kernel_size=3, stride=1, padding=1)
        #self.conv2 = nn.Conv3d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1)

        self.down1 = DoubleConv(in_channels, 64)
        self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.down2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        #self.down3 = DoubleConv(128, 256)
        #self.pool3 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        #self.down4 = DoubleConv(256, 512)
        #self.pool4 = nn.MaxPool3d(kernel_size=(1, 2, 2))

        #self.lstm1 = nn.LSTM(out_channels, hidden_channels, num_layers=num_layers, batch_first=True)
        #self.lstm2 = nn.LSTM(out_channels, hidden_channels, num_layers=num_layers, batch_first=True)
        #self.lstm3 = nn.LSTM(out_channels, hidden_channels, num_layers=num_layers, batch_first=True)
                
        self.lstm2 = crnn.Conv2dLSTM(in_channels=128,  
                                      out_channels=128,  
                                      kernel_size=self.kernel_size,  
                                      num_layers=self.num_layers, 
                                      bidirectional=True,
                                      dilation=self.dilation, 
                                      stride=self.stride, 
                                      dropout=self.dropout, 
                                      batch_first=True) 

        self.conv2 = nn.Conv3d(128*2, 128, kernel_size=1)

        self.up1 = nn.ConvTranspose3d(512, 256, kernel_size=(1, 2, 2), stride=(1, 2, 2))
        self.up_conv1 = DoubleConv(512, 256)
        self.up2 = nn.ConvTranspose3d(256, 128, kernel_size=(1, 2, 2), stride=(1, 2, 2))
        self.up_conv2 = DoubleConv(256, 128)
        self.up3 = nn.ConvTranspose3d(128, 64, kernel_size=(1, 2, 2), stride=(1, 2, 2), output_padding=(0, 1, 1))
        self.up_conv3 = DoubleConv(128, 64)
        self.out_conv = nn.Conv3d(64, out_channels, kernel_size=1)

    def forward(self, x):
        #x = self.conv1(x)
        #print(x.shape)
        #x = self.conv2(x)
         
        # Encoder
        x1 = self.down1(x)
        x2 = self.pool1(x1)
        x2 = self.down2(x2)
        #x3 = self.pool2(x2)
        #x3 = self.down3(x3)
        #x4 = self.pool3(x3)
        #x4 = self.down4(x4)

        # LSTM 

        x2 = x2.permute(0, 2, 1, 3, 4)
        x2, _ = self.lstm2(x2)
        x2 = x2.permute(0, 2, 1, 3, 4)
        x2 = self.conv2(x2)

        # Decoder
        #x = self.up1(x4)
        #x = torch.cat([x, x3], dim=1)
        #x = self.up_conv1(x)

        #x = self.up2(x)
        #x = torch.cat([x, x2], dim=1)
        #x = self.up_conv2(x)

        x = self.up3(x2)
        x = torch.cat([x, x1], dim=1)
        x = self.up_conv3(x)
        
        x = self.out_conv(x)        
        x = x.reshape(x.size(0), x.size(1), -1, x.size(3), x.size(4))
                
        #return x       
        return x[:, :, -48:, :, :] 
</code></pre>
<p>My images always have one channel (pixel values range from 0 to 255).<br>
By the way, do you see problems in this code I showed?</p>
<p>Thank you very much!</p>

J_Johnson: <p>I’m not at a computer, so unable to test, but will point a few things I noticed.</p>
<ol>
<li>Should be 2d UNet LSTM or 3d UNet. The 3d Unet eliminates the need for an LSTM since you’re feeding the entire time sequence at once(albeit more calculation intensive).</li>
<li>Self attention has good success in UNets. It helps the model focus on what’s important. Especially helps with skip connections, which I see you’ve included.</li>
</ol>
<p>Have you applied normalization between 0 and 1 for the inputs?</p>

marco_c: <p>Ok, considering I’m using 3D UNET and the input format I’m using is [batch_size=x, channels=1, size=144, width=7, height=7], would I be able to adapt my code to use 2D UNET to continue using LSTM?</p>
<p>Yes, I am normalizing the data into 0-1.</p>
<p>About self attention, thanks for the tip, I’ll add that.</p>
<p>Thank you very much!</p>

J_Johnson: <p>A 3D UNet or Vision Transformer would be superior to an RNN such as an LSTM.</p>
<p>Consider this, what is the highest accuracy you’ve seen a model attain? The LSTM has to learn what information to pass on and what to discard. Let’s suppose, after much training, it does so correctly 80% of the time. That is an 80% memory accuracy(which for an ML model is pretty good). That means by the 4th time sequence, you’ve lost upward of 59% (1-0.8^4) fidelity.</p>
<p>This is why language models before 2017 were not very good, compared to today. Why settle for 80% memory capture between time frames when you can just give a model the data at 100% fidelity?</p>
<p>By passing in all time frames, as with a 3D UNet, you no longer need that middle LSTM layer. You can pass the data directly from encoder to decoder.</p>

marco_c: <p>Hi! Right, I’m testing now the 3D Unet without the LSTM. I’m using this code, what do you think?</p>
<pre><code class="lang-auto">class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv3d(in_channels, out_channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv3d(out_channels, out_channels, kernel_size=(1, 3, 3), padding=(0, 1, 1)),
            nn.BatchNorm3d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)

class Unet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Unet, self).__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels

        self.down1 = DoubleConv(in_channels, 64)
        self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.down2 = DoubleConv(64, 128)
        self.pool2 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.conv2 = nn.Conv3d(128, 128, kernel_size=1)
        self.up3 = nn.ConvTranspose3d(128, 64, kernel_size=(1, 2, 2), stride=(1, 2, 2), output_padding=(0, 1, 1))
        self.up_conv3 = DoubleConv(128, 64)
        self.out_conv = nn.Conv3d(64, out_channels, kernel_size=1)

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.pool1(x1)
        x2 = self.down2(x2)
        x2 = self.conv2(x2)
        x = self.up3(x2)
        x = torch.cat([x, x1], dim=1)
        x = self.up_conv3(x)
        x = self.out_conv(x)
        x = x.reshape(x.size(0), x.size(1), -1, x.size(3), x.size(4))
        return x[:, :, -48:, :, :]

</code></pre>
<p>Thank you so much for all the information and help you have given me!</p>

J_Johnson: <p>That looks better.</p>
<p>I noticed you applied a skip connection near the beginning of the decoder but nowhere else. I’m guessing that was due to sizing issues since your output size is not the same as your input size. One way you can work around that is by just running a given skip tensor from the encoder through an nn.AdaptiveMaxPool3d or nn.AdaptiveAvgPool3d layer to get dim=2 to the right size needed. That will allow you to still pass those skip tensors for better fidelity when much of a scene remains the same in the output as it is in the input.</p>

marco_c: <p>Hi. Ok, nice! About the skip connection, I think you’re talking about this line:</p>
<pre><code class="lang-auto">x = torch.cat([x, x1], dim=1)
</code></pre>
<p>I’m doing this just to allow low-level information captured in the first layers of the network to be preserved and combined with high-level information in subsequent layers, to help improve the model’s ability to capture both fine detail and more abstract level features. But I’ll try to do it the way you suggested.</p>
<p>I made a version of the Vision Transformer network as below, what do you think? I haven’t tested it yet.</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
from torch.nn import Transformer

class VisionTransformer(nn.Module):
    def __init__(self, in_channels, out_channels, seq_length=144, embed_dim=256, num_heads=8, num_layers=6):
        super(VisionTransformer, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.seq_length = seq_length
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_layers = num_layers

        self.embedding = nn.Conv2d(in_channels, embed_dim, kernel_size=1)
        self.pos_embedding = nn.Parameter(torch.randn(1, seq_length, embed_dim))
        self.transformer = Transformer(
            d_model=embed_dim,
            nhead=num_heads,
            num_encoder_layers=num_layers,
            num_decoder_layers=num_layers
        )
        self.fc = nn.Linear(embed_dim, out_channels)
        self.output_conv = nn.Conv3d(out_channels, out_channels, kernel_size=1)

    def forward(self, x):
        bs, _, _, h, w = x.shape
        x = x.view(-1, self.in_channels, h, w)

        x = self.embedding(x)

        # Adjust pos_embedding size to match x
        pos_embedding = self.pos_embedding.repeat(bs, h, w, 1)

        x += pos_embedding

        x = x.permute(0, 2, 3, 1)
        x = x.view(bs, -1, self.embed_dim)

        x = self.transformer(x, x)

        x = x.view(bs, h, w, self.embed_dim)
        x = x.permute(0, 3, 1, 2)

        x = self.fc(x)

        x = x.view(bs, self.out_channels, -1, h, w)

        x = self.output_conv(x)

        x = x.view(bs, -1, h, w)

        return x
</code></pre>
<p>Thanks!</p>

J_Johnson: <p>Regarding ViTs, I’ve only seen them used well in classification type problems. Not in image generation. If you go that route, you may need to read up on a few papers with keywords ViT and image generation. From what I have read, you may need to add a positional encoder.</p>