DIM/image_encoder.py at master · TiSU32/DIM · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
'''Simple Image encoder mmdel.

'''

import cortex
from cortex.plugins import ModelPlugin
from cortex.built_ins.models.utils import update_encoder_args, update_decoder_args
from cortex.built_ins.models.classifier import SimpleClassifier, SimpleAttributeClassifier
from cortex.built_ins.networks.resnets import ResBlock
import torch


from alexnet_64 import AlexNetEncoder


class EncoderBase(ModelPlugin):

    def build(self, encoder_type: str='convnet', dim_out: int=None,
              encoder_args=dict(), semi_supervised=None):
        '''

        Args:
            encoder_type: Encoder model type.
            dim_out: Output size.
            encoder_args: Arguments for encoder build.

        '''
        x_shape = self.get_dims('x', 'y', 'c')
        if 'encoder' not in self.nets:
            if encoder_type == 'alexnet':
                Encoder = AlexNetEncoder
                encoder_args.pop('min_dim')
                encoder = Encoder(x_shape, dim_out=dim_out, **encoder_args)
            else:
                Encoder, _ = update_encoder_args(
                    x_shape, model_type=encoder_type)
                encoder = Encoder(x_shape, dim_out=dim_out, **encoder_args)

            self.nets.encoder = encoder

        self.data.reset(mode='test', make_pbar=False)
        self.data.next()
        X = self.inputs('inputs').cpu()
        Z = self.nets.encoder(X)
        dim_z = Z.size(1)

        self.linear_indices = []
        self.conv_indices = []
        self.c_idx = None

        index = -2
        try:
            while True:
                layer = self.nets.encoder.models[index]
                if isinstance(layer, torch.nn.modules.linear.Linear):
                    # Is a linear layer.
                    self.linear_indices.append(index)
                elif isinstance(layer, cortex.built_ins.networks.modules.View):
                    # Is a flattened version of a convolutional layer.
                    self.c_idx = self.c_idx or index
                elif isinstance(layer, torch.nn.modules.conv.Conv2d):
                    # Is a convolutional layer.
                    self.conv_indices.append(index)
                index -= 1
        except IndexError:
            pass

        if not semi_supervised:
            # Build the classifier on top of Y.
            self.classifier.build(dim_in=dim_z)

            # Build the classifier on top of the layer below Y.
            self.h_idx = self.linear_indices[0]
            dim_h = self.nets.encoder.states[self.h_idx].size(1)
            self.classifier_h.build(dim_in=dim_h)

        # Build the classifier on top of the last conv layer.
        dim_c = self.nets.encoder.states[self.c_idx].size(1)
        self.classifier_c.build(dim_in=dim_c)

    def encode(self, inputs, **kwargs):
        return self.nets.encoder(inputs, **kwargs)


class ImageEncoder(EncoderBase):
    '''Builds a simple image encoder.

    '''

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.classifier = SimpleClassifier(
            contract=dict(kwargs=dict(dim_in='dim_z')))
        self.classifier_h = SimpleClassifier(
            contract=dict(nets=dict(classifier='classifier_h'),
                          kwargs=dict(classifier_args='classifier_h_args')))
        self.classifier_c = SimpleClassifier(
            contract=dict(nets=dict(classifier='classifier_c'),
                          kwargs=dict(classifier_args='classifier_c_args')))

    def routine(self, inputs, targets, semi_supervised=False):
        '''

        Args:
            semi_supervised: If set, pass classification gradients through
                encoder (off by default).

        '''
        Z_Q = self.encode(inputs, nonlinearity=False).detach()

        if semi_supervised:
            self.classifier_c.routine(
                self.nets.encoder.states[self.c_idx], targets)
            if 'classifier_c' in self.losses:
                # For STL, some batches will not have labels, so no loss.
                if 'encoder' in self.losses:
                    self.losses.encoder += self.losses.classifier_c
                else:
                    self.losses.encoder = self.losses.classifier_c

        else:
            self.classifier.routine(Z_Q.detach(), targets)
            self.classifier_h.routine(
                self.nets.encoder.states[self.h_idx].detach(), targets)
            self.classifier_c.routine(
                self.nets.encoder.states[self.c_idx].detach(), targets)