def __init__(self,
num_classes,
backbone,
backbone_indices,
n_scales=[0.5, 1.0, 2.0],
ocr_mid_channels=512,
ocr_key_channels=256,
align_corners=False,
pretrained=None):
super().__init__()
self.ocrnet = OCRNetNV(num_classes,
backbone,
backbone_indices,
ocr_mid_channels=ocr_mid_channels,
ocr_key_channels=ocr_key_channels,
align_corners=align_corners,
ms_attention=True)
self.scale_attn = AttenHead(in_ch=ocr_mid_channels, out_ch=1)
self.n_scales = n_scales
self.pretrained = pretrained
self.align_corners = align_corners
if self.pretrained is not None:
utils.load_pretrained_model(self, self.pretrained)
def init_weight(self):
for layer in self.sublayers():
if isinstance(layer, nn.Conv2D):
param_init.normal_init(layer.weight, std=0.001)
elif isinstance(layer, (nn.BatchNorm, nn.SyncBatchNorm)):
param_init.constant_init(layer.weight, value=1.0)
param_init.constant_init(layer.bias, value=0.0)
if self.pretrained is not None:
utils.load_pretrained_model(self, self.pretrained)
def init_weights(self, pretrained=None):
"""Initialize the weights in backbone.
Args:
pretrained (str, optional): Path to pre-trained weights.
Defaults to None.
"""
if pretrained is not None:
utils.load_pretrained_model(self, self.pretrained)
else:
for sublayer in self.sublayers():
if isinstance(sublayer, nn.Linear):
trunc_normal_(sublayer.weight)
if isinstance(sublayer,
nn.Linear) and sublayer.bias is not None:
zeros_(sublayer.bias)
elif isinstance(sublayer, nn.LayerNorm):
zeros_(sublayer.bias)
ones_(sublayer.weight)
def init_weight(self):
utils.load_pretrained_model(self, self.pretrained)
# load and resize pos_embed
model_path = self.pretrained
if not os.path.exists(model_path):
model_path = utils.download_pretrained_model(model_path)
load_state_dict = paddle.load(model_path)
model_state_dict = self.state_dict()
pos_embed_name = "pos_embed"
if pos_embed_name in load_state_dict.keys():
load_pos_embed = paddle.to_tensor(load_state_dict[pos_embed_name],
dtype="float32")
if self.pos_embed.shape != load_pos_embed.shape:
pos_size = int(math.sqrt(load_pos_embed.shape[1] - 1))
model_state_dict[pos_embed_name] = self.resize_pos_embed(
load_pos_embed, (pos_size, pos_size),
(self.pos_h, self.pos_w))
self.set_dict(model_state_dict)
logger.info(
"Load pos_embed and resize it from {} to {} .".format(
load_pos_embed.shape, self.pos_embed.shape))
def init_weight(self):
if self.pretrained is not None:
utils.load_pretrained_model(self, self.pretrained)
def init_weight(self):
utils.load_pretrained_model(self, self.pretrained)
def init_weight(self):
if self.pretrained is not None:
utils.load_pretrained_model(self, self.pretrained)
else:
self.apply(self._init_weights)
