Method forward

modules/VIPTRv2.py:841–914

(self, x, size)

            self.downsample = None

    def forward(self, x, size):
        b, h, w, d = x.size()
        # print(x.size())
        if self.mixer_type == 'Local1':
            for blk in self.blocks:
                x = x.flatten(1).reshape(b, -1, d)
                x = blk(x, size)
                x = x.reshape(b, h, w, -1)

        elif self.mixer_type == 'Local2':
            rel_pos = self.Relpos((h, w), chunkwise_recurrent=self.chunkwise_recurrent)
            for blk in self.blocks:
                if self.use_checkpoint:
                    tmp_blk = partial(blk, incremental_state=None, chunkwise_recurrent=self.chunkwise_recurrent,
                                      retention_rel_pos=rel_pos)
                    x = checkpoint.checkpoint(tmp_blk, x)
                else:
                    x = blk(x, incremental_state=None, chunkwise_recurrent=self.chunkwise_recurrent,
                            retention_rel_pos=rel_pos)

        elif self.mixer_type == 'Global1':
            for blk in self.blocks:
                x = x.flatten(1).reshape(b, -1, d)
                # x = x.permute(0, 3, 1, 2).contiguous()
                x = blk(x, size)
                x = x.reshape(b, h, w, -1)

        elif self.mixer_type == 'Global2':
            for blk in self.blocks:
                x = x.flatten(1).reshape(b, -1, d)
                x = blk(x, size)
                x = x.reshape(b, h, w, -1)

        elif self.mixer_type == 'LG1':
            # print(self.mixer_type)
            # print(x.shape)
            # rel_pos = self.Relpos((h, w), chunkwise_recurrent=self.chunkwise_recurrent)
            for lblk, gblk in zip(self.local_unit, self.global_unit):
                x = x.flatten(1).reshape(b, -1, d)
                x1, x2 = torch.chunk(x, chunks=2, dim=2)
                x1 = lblk(x1, size)
                x2 = gblk(x2, size)
                x = torch.cat([x1, x2], dim=2)
                x = x.transpose(1, 2).contiguous().reshape(b, -1, h, w)
                x = self.proj(x) + x
                x = x.permute(0, 2, 3, 1).contiguous()
                x = x.reshape(b, h, w, -1)

                # x = x.flatten(1).reshape(b, -1, d)
                # x1, x2 = torch.chunk(x, chunks=2, dim=2)
                # x1 = lblk(x1, size)
                # x2 = gblk(x2, size)
                # x = torch.cat([x1, x2], dim=2)
                # x = x.transpose(1, 2).contiguous().reshape(b, -1, h, w)
                # x = self.proj(x) + x
                # x = x.permute(0, 2, 3, 1).contiguous()
                # x = x.reshape(b, h, w, -1)