Source code for openhands.models.encoder.cnn2d

import torch
import torch.nn as nn
import timm


class CNN2D(nn.Module):
    """
    Creates a 2D Convolution backbone from `timm` library 
    
    Args:
        in_channels (int): Number of input channels
        backbone (string): Backbone to use
        pretrained (bool, optional): Whether to use pretrained Backbone. Default: ``True``
    """
    def __init__(self, in_channels=3, backbone="resnet18", pretrained=True):
        super().__init__()
        assert in_channels == 3
        self.backbone = timm.create_model(backbone, pretrained=pretrained)
        self.n_out_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Identity()

    def forward(self, x):
        """
        forward step
        """
        b, c, t, h, w = x.shape
        cnn_embeds = []
        for i in range(t):
            out = self.backbone(x[:, :, i, :, :])
            out = out.view(out.shape[0], -1)
            cnn_embeds.append(out)

        return torch.stack(cnn_embeds, dim=0).transpose(0, 1) # Batch-first