src/transformers/models/flava/modeling_flava.py/FlavaImageCodebook/__init__
class FlavaImageCodebook:    def __init__(
        self,
        config: FlavaImageCodebookConfig,
        **kwargs: Any,
    ):
        super().__init__(config)

        self.config = config
        self.num_groups = config.num_groups
        self.input_channels = config.input_channels
        self.num_blocks_per_group = config.num_blocks_per_group
        self.hidden_size = config.hidden_size
        self.vocab_size = config.vocab_size

        num_layers = self.num_groups * self.num_blocks_per_group

        output_blocks = OrderedDict()
        output_blocks["relu"] = nn.ReLU()
        output_blocks["conv"] = nn.Conv2d(8 * self.hidden_size, self.vocab_size, kernel_size=1, padding=0)

        blocks = OrderedDict()
        blocks["input"] = nn.Conv2d(self.input_channels, 1 * self.hidden_size, kernel_size=7, padding=3)
        blocks["group_1"] = FlavaImageCodebookLayerGroup(
            self.num_blocks_per_group, num_layers, 1 * self.hidden_size, 1 * self.hidden_size
        )
        blocks["group_2"] = FlavaImageCodebookLayerGroup(
            self.num_blocks_per_group, num_layers, 1 * self.hidden_size, 2 * self.hidden_size
        )
        blocks["group_3"] = FlavaImageCodebookLayerGroup(
            self.num_blocks_per_group, num_layers, 2 * self.hidden_size, 4 * self.hidden_size
        )
        blocks["group_4"] = FlavaImageCodebookLayerGroup(
            self.num_blocks_per_group, num_layers, 4 * self.hidden_size, 8 * self.hidden_size, use_pool=False
        )
        blocks["output"] = nn.Sequential(output_blocks)

        self.blocks = nn.Sequential(blocks)

        self.post_init()

        if self.config.freeze:
            for param in self.parameters():
                param.requires_grad = False