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())
