pytorch - 💡(How to fix) Fix [torch.compile] Graph break and assertion error when using getattr with None default followed by call [1 comments, 2 participants]

AssertionError: expected size 3==3, stride 1==49 at dim=1

---

import os
import torch
import torch.nn as nn
import torch.nn.functional as F

class MaskedConv2d(nn.Conv2d):

    def __init__(self, mask_type, *args, **kwargs):
        super(MaskedConv2d, self).__init__(*args, **kwargs)
        assert mask_type in ['A', 'B'], "mask_type must be either 'A' or 'B'"
        self.register_buffer('mask', self.weight.data.clone())
        (_, _, height, width) = self.weight.size()
        self.mask.fill_(1)
        self.mask[:, :, height // 2, width // 2 + (mask_type == 'B'):] = 0
        self.mask[:, :, height // 2 + 1:] = 0

    def forward(self, x):
        self.weight.data *= self.mask
        return super(MaskedConv2d, self).forward(x)

class ResidualBlock(nn.Module):

    def __init__(self, in_channels):
        super(ResidualBlock, self).__init__()
        self.conv = nn.Sequential(nn.ReLU(), MaskedConv2d('B', in_channels, in_channels // 2, 1, 1, 0, bias=False), nn.BatchNorm2d(in_channels // 2), nn.ReLU(), MaskedConv2d('B', in_channels // 2, in_channels // 2, 3, 1, 1, bias=False), nn.BatchNorm2d(in_channels // 2), nn.ReLU(), MaskedConv2d('B', in_channels // 2, in_channels, 1, 1, 0, bias=False), nn.BatchNorm2d(in_channels))
        self.skip = nn.Identity()

    def forward(self, x):
        return self.conv(x) + self.skip(x)

class PixelCNN(nn.Module):

    def __init__(self, in_channels=3, hidden_dim=128, out_channels=256, num_layers=12):
        super(PixelCNN, self).__init__()
        self.input_conv = MaskedConv2d('A', in_channels, hidden_dim, 7, 1, 3)
        self.res_blocks = nn.ModuleList()
        for _ in range(num_layers):
            self.res_blocks.append(ResidualBlock(hidden_dim))
        self.output_layers = nn.Sequential(nn.ReLU(), MaskedConv2d('B', hidden_dim, hidden_dim, 1), nn.ReLU(), MaskedConv2d('B', hidden_dim, out_channels * in_channels, 1))
        self.out_channels = out_channels
        self.in_channels = in_channels

    def forward(self, x):
        x = getattr(x, 'float', None)() / 255.0
        h = self.input_conv(x)
        for res_block in self.res_blocks:
            h = res_block(h)
        out = self.output_layers(h)
        (batch_size, _, height, width) = out.shape
        out = out.view(batch_size, self.out_channels, self.in_channels, height, width)
        return out

    def sample(self, shape, device='cpu', temperature=1.0):
        with torch.no_grad():
            samples = torch.zeros(shape, device=device)
            for i in range(shape[2]):
                for j in range(shape[3]):
                    for c in range(shape[1]):
                        logits = self.forward(samples * 255.0)[:, :, c, i, j] / temperature
                        probs = F.softmax(logits, dim=1)
                        sampled = torch.multinomial(probs, 1).squeeze(-1)
                        samples[:, c, i, j] = sampled.float() / 255.0
            return samples * 255.0

def get_default_model():
    in_channels = 3
    hidden_dim = 128
    out_channels = 256
    num_layers = 12
    model = PixelCNN(in_channels=in_channels, hidden_dim=hidden_dim, out_channels=out_channels, num_layers=num_layers)
    return model

def get_sample_inputs():
    batch_size = 2
    channels = 3
    height = 32
    width = 32
    x = torch.randint(0, 256, (batch_size, channels, height, width), dtype=torch.float32)
    return (x,)

def main():
    model = get_default_model()
    model.eval()
    inputs = get_sample_inputs()
    with torch.no_grad():
        output = model(*inputs)
    print('Model executed successfully!')
    print(f'Input shape: {inputs[0].shape}')
    print(f'Output shape: {output.shape}')
    print(f'Model parameters: {sum((p.numel() for p in model.parameters()))}')
    compiled_model = torch.compile(model)
    with torch.no_grad():
        output_compile = compiled_model(*inputs)
    print(f'Compile  shape: {output_compile.shape}')
if __name__ == '__main__':
    main()

---

AssertionError: expected size 3==3, stride 1==49 at dim=1