def create_conv_layers(self, shape, conv_args): '''Creates a set of convolutional layers. Args: shape: Input shape. conv_args: List of tuple of convolutional arguments. Returns: nn.Sequential: a sequence of convolutional layers. ''' conv_layers = nn.Sequential() conv_args = conv_args or [] dim_x, dim_y, dim_in = shape for i, (dim_out, f, s, p, batch_norm, dropout, nonlinearity, pool) in enumerate(conv_args): name = '({}/{})_{}'.format(dim_in, dim_out, i + 1) conv_block = nn.Sequential() if dim_out is not None: conv = nn.Conv2d(dim_in, dim_out, kernel_size=f, stride=s, padding=p, bias=not (batch_norm)) conv_block.add_module(name + 'conv', conv) dim_x, dim_y = self.next_size(dim_x, dim_y, f, s, p) else: dim_out = dim_in if dropout: conv_block.add_module(name + 'do', nn.Dropout2d(p=dropout)) if batch_norm: bn = nn.BatchNorm2d(dim_out) conv_block.add_module(name + 'bn', bn) if nonlinearity: nonlinearity = get_nonlinearity(nonlinearity) conv_block.add_module(nonlinearity.__class__.__name__, nonlinearity) if pool: (pool_type, kernel, stride) = pool Pool = getattr(nn, pool_type) conv_block.add_module(name + 'pool', Pool(kernel_size=kernel, stride=stride)) dim_x, dim_y = self.next_size(dim_x, dim_y, kernel, stride, 0) conv_layers.add_module(name, conv_block) dim_in = dim_out dim_out = dim_in return conv_layers, (dim_x, dim_y, dim_out)
def test_get_nonlinearity(nonlinearity): """ Args: nonlinearity(@pytest.fixture): dict Asserts: True if right instance of activation function is returned. """ relu = get_nonlinearity(nonlinearity['relu']) tanh = get_nonlinearity(nonlinearity['tanh']) leakyrelu = get_nonlinearity(nonlinearity['leakyrelu']) sigmoid = get_nonlinearity(nonlinearity['sigmoid']) assert callable(sigmoid) assert callable(tanh) assert isinstance(relu, nn.modules.activation.ReLU) assert isinstance(leakyrelu, nn.modules.activation.LeakyReLU)
def finish_block(self, block, shape, bn=False, ln=False, do=False, act=None, pool=None): '''Finishes a block. Adds batch norm, dropout, activation, pooling. Args: block (nn.Sequential): Block to add conv layer to. shape (tuple): Shape of the input. bn (bool): Batch normalization. ln (bool): Layer normalization. do (float): Dropout. act (str): Activation. pool (tuple): Pooling. In format (pool type, kernel size, stride). Returns: ''' if len(shape) == 1: BN = nn.BatchNorm1d DO = nn.Dropout elif len(shape) == 3: BN = nn.BatchNorm2d DO = nn.Dropout2d else: raise NotImplementedError('Shape {} not supported'.format(shape)) LN = nn.LayerNorm if ln and bn: raise ValueError('Use only one sort of normalization.') dim_out = shape[-1] if do: block.add_module('do', DO(p=do)) if bn: block.add_module('bn', BN(dim_out)) if ln: block.add_module('ln', LN(dim_out)) if act: nonlinearity = get_nonlinearity(act) block.add_module(nonlinearity.__class__.__name__, nonlinearity) if pool: if len(shape) == 1: raise ValueError('Cannot pool on 1d tensor.') (pool_type, kernel, stride) = pool Pool = getattr(nn, pool_type) block.add_module('pool', Pool(kernel_size=kernel, stride=stride)) dim_x, dim_y, dim_out = shape dim_x, dim_y = self.next_conv_size(dim_x, dim_y, kernel, stride, 0) shape = (dim_x, dim_y, dim_out) return shape
def create_linear_layers(self, dim_in, fc_args): ''' Args: dim_in: Number of input units. fc_args: List of tuple of fully-connected arguments. Returns: nn.Sequential. ''' fc_layers = nn.Sequential() fc_args = fc_args or [] for i, (dim_out, batch_norm, dropout, nonlinearity) in enumerate(fc_args): name = '({}/{})_{}'.format(dim_in, dim_out, i + 1) fc_block = nn.Sequential() if dim_out is not None: fc_block.add_module(name + 'fc', nn.Linear(dim_in, dim_out)) else: dim_out = dim_in if dropout: fc_block.add_module(name + 'do', nn.Dropout(p=dropout)) if batch_norm: bn = nn.BatchNorm1d(dim_out) fc_block.add_module(name + 'bn', bn) if nonlinearity: nonlinearity = get_nonlinearity(nonlinearity) fc_block.add_module(nonlinearity.__class__.__name__, nonlinearity) fc_layers.add_module(name, fc_block) dim_in = dim_out return fc_layers, dim_in