Entroformer图像编码

class TransDecoder2(TransDecoder):
    train_scan_mode = 'default'  #  'random', 'default'
    test_scan_mode = 'checkboard'
    is_decoder = False
    def __init__(self, cin=0, cout=0, opt=None):
        super().__init__(cin, cout, opt)
        del self.sos_pred_token

    def forward(self, x, manual_mask=None):
        x = x.clone()
        batch_size, channels, height, width  = x.shape   # input_shape

        # Self-attention Mask & Token Mask
        if manual_mask is None:
            mask, token_mask, input_mask, output_mask = self.get_mask(batch_size, height, width)
        else:
            mask, token_mask, input_mask, output_mask = manual_mask
        mask, input_mask, output_mask = mask.to(x.device), input_mask.to(x.device), output_mask.to(x.device)
        token_mask = token_mask.to(x.device) if token_mask is not None else token_mask

        # Mask Input
        x.masked_fill_(~input_mask, 0.)
        # Input Embedding
        x = rearrange(x, 'b c h w -> b (h w) c')
        inputs_embeds = self.to_patch_embedding(x)

        # Init state
        position_bias = None
        # encoder_decoder_position_bias = None
        hidden_states = inputs_embeds

        topk = self.attn_topk
        if self.training and topk != -1:
            topk = np.random.randint(topk//2, topk*2)

        for _, layer_module in enumerate(self.blocks):
            # Transformer block
            layer_outputs = layer_module(
                hidden_states,
                shape_2d=[height, width],
                attention_mask=mask,
                position_bias=position_bias,
                topk=topk,
            )

            hidden_states = layer_outputs[0]

            # We share the position biases between the layers - the first layer store them
            # layer_outputs = hidden-states, (self-attention position bias), (cross-attention position bias)
            if self.rpe_shared:
                position_bias = layer_outputs[1]

        # Out projection
        out = self.mlp_head(hidden_states)        
        # Reshape Output to 2D map
        out = rearrange(out, 'b (h w) c -> b c h w', h=height)
        # Mask output
        out.masked_fill_(~output_mask, 0.)
        return out            

    def get_mask(self, b, h, w):
        n = h*w
        if self.training:
            if(self.train_scan_mode == 'random' and hasattr(self, 'sampler')):
                #mask = torch.ones(n, n).bool()    # modified
                token_mask = None
                mask_random = (self.sampler.sample([n]) > self.mask_ratio).bool()
                input_mask = mask_random.clone().view(h,w)
                output_mask = ~mask_random.clone().view(h,w)
                mask = repeat(mask_random.unsqueeze(0), '() n -> d n', d=n)
                mask = mask | torch.eye(n).bool()
            else:
                #mask = torch.ones(n, n).bool()
                token_mask = None
                mask_checkboard = torch.ones((h, w)).bool()
                mask_checkboard[0::2, 0::2] = 0
                mask_checkboard[1::2, 1::2] = 0
                input_mask = mask_checkboard.clone()
                output_mask = ~mask_checkboard.clone()
                mask = repeat(mask_checkboard.view(1,-1), '() n -> d n', d=n)
                mask = mask | torch.eye(n).bool()
        else:
            if 'checkboard' in self.test_scan_mode:
                #mask = torch.ones(n, n).bool()
                token_mask = None
                mask_checkboard = torch.ones((h, w)).bool()
                if self.test_scan_mode == 'checkboard':
                    mask_checkboard[0::2, 0::2] = 0
                    mask_checkboard[1::2, 1::2] = 0
                else:
                    mask_checkboard[0::2, 1::2] = 0
                    mask_checkboard[1::2, 0::2] = 0
                input_mask = mask_checkboard.clone()
                output_mask = ~mask_checkboard.clone()
                mask = repeat(mask_checkboard.view(1,-1), '() n -> d n', d=n)
                mask = mask | torch.eye(n).bool()
            else:
                raise ValueError("No such test scan mode.")

        #print(input_mask)
        mask = repeat(mask.unsqueeze(0), '() d n -> b d n', b=b)
        token_mask = token_mask  # torch.ones_like(mask).bool()
        input_mask = repeat(input_mask.unsqueeze(0).unsqueeze(0), '() () h w -> b d h w', b=b, d=self.cin)
        channel = self.dim if self.cout == 0 else self.cout
        output_mask = repeat(output_mask.unsqueeze(0).unsqueeze(0), '() () h w -> b d h w', b=b, d=channel)

        return mask, token_mask, input_mask, output_mask

sh scripts/pretrain.sh 0.3
sh scripts/train.sh [tradeoff_lambda(e.g. 0.02)]
(You may use your own dataset by modifying the train/test data path.)

# Kodak
sh scripts/test.sh [/path/to/kodak] [model_path]
(sh test_parallel.sh [/path/to/kodak] [model_path])

sh scripts/compress.sh [original.png] [model_path]
(sh compress_parallel.sh [original.png] [model_path])

sh scripts/decompress.sh [original.bin] [model_path]
(sh decompress_parallel.sh [original.bin] [model_path])