def __init__(self, cell, generator, num_steps, num_class, input_attention_block=None, output_attention_block=None, text_transform=None, holistic_input_from=None): super(AttHead, self).__init__() # from vedastr import utils # utils.set_random_seed(1) if input_attention_block is not None: self.input_attention_block = build_brick(input_attention_block) self.cell = build_sequence_decoder(cell) self.generator = build_torch_nn(generator) self.num_steps = num_steps self.num_class = num_class if output_attention_block is not None: self.output_attention_block = build_brick(output_attention_block) if text_transform is not None: self.text_transform = build_torch_nn(text_transform) if holistic_input_from is not None: self.holistic_input_from = holistic_input_from self.register_buffer('embeddings', torch.diag(torch.ones(self.num_class))) logger.info('AttHead init weights') init_weights(self.modules())
def __init__(self, attention, attention_norm, feedforward, feedforward_norm): super(TransformerEncoderLayer2D, self).__init__() self.attention = build_attention(attention) self.attention_norm = build_torch_nn(attention_norm) self.feedforward = build_feedforward(feedforward) self.feedforward_norm = build_torch_nn(feedforward_norm)
def __init__(self, input_pool, layers, keep_order=False): super(RNN, self).__init__() self.keep_order = keep_order if input_pool: self.input_pool = build_torch_nn(input_pool) self.layers = nn.ModuleList() for i, (layer_name, layer_cfg) in enumerate(layers): if layer_name in ['rnn', 'fc']: self.layers.add_module('{}_{}'.format(i, layer_name), build_torch_nn(layer_cfg)) else: raise ValueError('Unknown layer name {}'.format(layer_name)) init_weights(self.modules())
def __init__(self, F, input_size, output_size, stn): super(TPS_STN, self).__init__() self.F = F self.input_size = input_size self.output_size = output_size self.feature_extractor = build_feature_extractor(stn['feature_extractor']) self.pool = build_torch_nn(stn['pool']) heads = [] for head in stn['head']: heads.append(build_module(head)) self.heads = nn.Sequential(*heads) self.grid_generator = GridGenerator(F, output_size) # Init last fc in heads last_fc = heads[-1].fc last_fc.weight.data.fill_(0) """ see RARE paper Fig. 6 (a) """ ctrl_pts_x = np.linspace(-1.0, 1.0, int(F / 2)) ctrl_pts_y_top = np.linspace(0.0, -1.0, num=int(F / 2)) ctrl_pts_y_bottom = np.linspace(1.0, 0.0, num=int(F / 2)) ctrl_pts_top = np.stack([ctrl_pts_x, ctrl_pts_y_top], axis=1) ctrl_pts_bottom = np.stack([ctrl_pts_x, ctrl_pts_y_bottom], axis=1) initial_bias = np.concatenate([ctrl_pts_top, ctrl_pts_bottom], axis=0) last_fc.bias.data = torch.from_numpy(initial_bias).float().view(-1)
def __init__(self, in_channels, out_channels, num_class, batch_max_length, from_layer, inner_channels=None, bias=True, activation='relu', inplace=True, dropouts=None, num_fcs=0, pool=None): super(FCHead, self).__init__() self.num_class = num_class self.batch_max_length = batch_max_length self.from_layer = from_layer if num_fcs > 0: inter_fc = FCModules(in_channels, inner_channels, bias, activation, inplace, dropouts, num_fcs) fc = nn.Linear(inner_channels, out_channels) else: inter_fc = nn.Sequential() fc = nn.Linear(in_channels, out_channels) if pool is not None: self.pool = build_torch_nn(pool) self.inter_fc = inter_fc self.fc = fc logger.info('FCHead init weights') init_weights(self.modules())
def __init__(self, layers, zero_init_residual=False, groups=1, width_per_group=64, norm_layer=None): super(GResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.dilation = 1 self.groups = groups self.base_width = width_per_group self.layers = nn.ModuleList() stage_layers = [] for layer_name, layer_cfg in layers: if layer_name == 'conv': layer = build_module(layer_cfg) self.inplanes = layer_cfg['out_channels'] elif layer_name == 'pool': layer = build_torch_nn(layer_cfg) elif layer_name == 'block': layer = self._make_layer(**layer_cfg) else: raise ValueError('Unknown layer name {}'.format(layer_name)) stride = layer_cfg.get('stride', 1) max_stride = stride if isinstance(stride, int) else max(stride) if max_stride > 1: self.layers.append(nn.Sequential(*stage_layers)) stage_layers = [] stage_layers.append(layer) self.layers.append(nn.Sequential(*stage_layers)) logger.info('GResNet init weights') init_weights(self.modules())
def __init__(self, cfg): super(SPN, self).__init__() self.body = build_feature_extractor(cfg['feature_extractor']) self.pool = build_torch_nn(cfg['pool']) heads = [] for head in cfg['head']: heads.append(build_module(head)) self.head = nn.Sequential(*heads)
def __init__(self, layers, zero_init_residual=False, groups=1, width_per_group=64, norm_layer=None): super(GResNet, self).__init__() if norm_layer is None: norm_layer = nn.BatchNorm2d self._norm_layer = norm_layer self.dilation = 1 self.groups = groups self.base_width = width_per_group self.layers = nn.ModuleList() stage_layers = [] for layer_name, layer_cfg in layers: if layer_name == 'conv': layer = build_module(layer_cfg) self.inplanes = layer_cfg['out_channels'] elif layer_name == 'pool': layer = build_torch_nn(layer_cfg) elif layer_name == 'block': layer = self._make_layer(**layer_cfg) else: raise ValueError('Unknown layer name {}'.format(layer_name)) stride = layer_cfg.get('stride', 1) max_stride = stride if isinstance(stride, int) else max(stride) if max_stride > 1: self.layers.append(nn.Sequential(*stage_layers)) stage_layers = [] stage_layers.append(layer) self.layers.append(nn.Sequential(*stage_layers)) for m in self.modules(): if isinstance(m, nn.Conv2d): nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu') elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): nn.init.constant_(m.weight, 1) nn.init.constant_(m.bias, 0) # Zero-initialize the last BN in each residual branch, # so that the residual branch starts with zeros, and each residual block behaves like an identity. # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 if zero_init_residual: for m in self.modules(): if isinstance(m, Bottleneck): nn.init.constant_(m.bn3.weight, 0) elif isinstance(m, BasicBlock): nn.init.constant_(m.bn2.weight, 0)
def __init__( self, decoder, generator, embedding, num_steps, pad_id, src_from, src_mask_from=None, ): super(TransformerHead, self).__init__() self.decoder = build_sequence_decoder(decoder) self.generator = build_torch_nn(generator) self.embedding = build_torch_nn(embedding) self.num_steps = num_steps self.pad_id = pad_id self.src_from = src_from self.src_mask_from = src_mask_from logger.info('TransformerHead init weights') init_weights(self.modules())
def __init__(self, layers): super(GVGG, self).__init__() self.layers = nn.ModuleList() stage_layers = [] for layer_name, layer_cfg in layers: if layer_name == 'conv': layer = build_module(layer_cfg) elif layer_name == 'pool': layer = build_torch_nn(layer_cfg) else: raise ValueError('Unknown layer name {}'.format(layer_name)) stride = layer_cfg.get('stride', 1) max_stride = stride if isinstance(stride, int) else max(stride) if max_stride > 1: self.layers.append(nn.Sequential(*stage_layers)) stage_layers = [] stage_layers.append(layer) self.layers.append(nn.Sequential(*stage_layers)) init_weights(self.modules())