def __init__()

in conv_lstm_models.py [0:0]

• 48 lines of code
• 3 McCabe index (conditional complexity)

    def __init__(self, args,
                 model_name='lstm',
                 map_embsize=64,
                 race_embsize=8,
                 residual=False,
                 dec_convsize=3,
                 dec_depth=3,
                 dec_embsize=128,
                 conv=nn.Conv2d,
                 nonlin=nn.ELU,
                 hid_dim=256,
                 rnn_input_size=2048,
                 lstm_nlayers=3,
                 lstm_dropout=0):
        super(lstm, self).__init__()
        self.model_name = model_name
        self.nonlin = nonlin(inplace=True) if nonlin is not None else IdentityFn()
        self.convmod = residual_layer(conv, self.nonlin) if residual else conv

        self.residual = residual
        self.hid_dim = hid_dim
        self.rnn_input_size = rnn_input_size
        self.lstm_num_layers = lstm_nlayers
        self.lstm_dropout = lstm_dropout
        self.dec_convsize = dec_convsize
        self.dec_depth = dec_depth
        self.dec_embsize = dec_embsize
        self.map_embsize = map_embsize
        self.race_embsize = race_embsize

        # with striding of 64, that's 16x16 on biggest maps (1024x1024),
        # thus is we reshape the map in a 1D vector that's 8 dims per (x,y)
        self.inp_embsize = self.rnn_input_size // (1024//args.stride)**2
        self.nfeat = args.n_inp_feats * 2
        self.nchannel = self.nfeat + race_embsize * 2 + map_embsize

        # Modules
        self.trunk = _MapRaceFeaturize(args, map_embsize, race_embsize)
        self.conv1x1 = nn.Conv2d(self.nchannel, self.inp_embsize, 1)  # TODO do that before trunk?
        self.rnn = nn.LSTM(self.rnn_input_size, self.hid_dim, self.lstm_num_layers, dropout=self.lstm_dropout)
        self.hidden = self.init_hidden()
        self.decoder = decoder(dec_convsize, dec_convsize)(
            conv        =conv,
            non_lin     =self.nonlin,
            input_size  =self.nchannel + self.hid_dim,
            interm_size =self.dec_embsize,
            output_size =self.dec_embsize,
            depth       =dec_depth
        ) # should be depth=1, the lstm should do work the work.
        self.regression_head = nn.Conv2d(self.dec_embsize, self.nfeat, 1)
        self.unit_class_head = nn.Conv2d(self.dec_embsize, 2 * len(utm.our_units_inds), 1)
        self.bldg_class_head = nn.Conv2d(self.dec_embsize, 2 * len(utm.our_bldgs_inds), 1)
        self.opbt_class_head = nn.Linear(self.dec_embsize, len(utm.nmy_bldgs_inds))
        self.accepts_bptt = True