intuitive262/adaptive_image_compression

1import torch2import torch.nn as nn3from utils import CNNBlock, TCNNBlock, _create_CNN_block4 5architecture = [6    (2,0,2,128, "U"),7    [(3,1,1,128, "C"), 2],8    (2,0,2,64, "U"),9    [(3,1,1,64, "C"), 2],10    (2,0,2,32, "U"),11    [(3,1,1,32, "C"), 2],12    (3,1,1,3, "C")13]14 15class Decoder(nn.Module):16    def __init__(self, in_channels=128, **kwargs):17        super(Decoder, self).__init__()18        self.in_channels = in_channels19        self.architecture = architecture20        # self.out_channels = out_channels21        22        self.decoder = self._create_decoder(self.in_channels, self.architecture)23 24    def forward(self, x):25        return self.decoder(x)26 27    def _create_decoder(self, in_channels, architecture):28        layers = []29 30        for archi_x in architecture:31            if type(archi_x) == list:32                blk_cfg, num_repeats = archi_x33                for _  in range(num_repeats):34                    layers.append(_create_CNN_block(35                        in_channels, blk_cfg36                    ))37                    in_channels = blk_cfg[-2]38            elif type(archi_x) == tuple:39                layers.append(_create_CNN_block(40                    in_channels, archi_x41                ))42                in_channels = archi_x[-2]43 44        return nn.Sequential(*layers)45    46    def test(self, x):47        y_ = self.forward(x)48        print(f"Decoder-Model(test for {x.shape}):",y_.shape)49    50def test(x):51    decoder_model = Decoder()52    y_ = decoder_model(x)53    print("\noutput_head:", y_.shape)54 55 56if __name__ == "__main__":57    test(torch.zeros((1,128,256,256)))