class LightHeadRCNNResNet101_Head(Module):
def __init__(self,
n_class=81,
roi_size=7,
out_channels=10,
spatial_scale=1 / 16.,
sampling_ratio=2,
roi_align=False):
super(LightHeadRCNNResNet101_Head, self).__init__()
self.n_class = n_class
self.spatial_scale = spatial_scale
self.roi_size = roi_size
self.out_channels = out_channels
self.sampling_ratio = sampling_ratio
self.c_out = self.roi_size * self.roi_size * self.out_channels
self.global_context_module = GlobalContextModule(
2048, 256, self.c_out, 15)
self.roi_align = roi_align
self.flatten = Flatten()
self.fc1 = Linear(self.c_out, 2048)
self.score = Linear(2048, n_class)
self.cls_loc = Linear(2048, 4 * n_class)
self.apply(lambda x: normal_init(x, 0, 0.01))
self.cls_loc.apply(lambda x: normal_init(x, 0, 0.001))
if self.roi_align:
self.pooling = RoIAlignMax(self.roi_size, self.roi_size,
self.spatial_scale, self.sampling_ratio)
# self.pooling = RoIAlign(self.roi_size, self.spatial_scale, self.sampling_ratio)
self.conv_pool = Conv2d(self.c_out,
self.out_channels,
1,
bias=False)
else:
self.pooling = PSRoIAlign(self.spatial_scale, self.roi_size,
self.sampling_ratio, self.out_channels)
def __call__(self, x, rois):
# global context module
device = x.device
h = self.global_context_module(x)
if self.roi_align:
func_roi = torch.cat(
(torch.zeros([rois.shape[0], 1],
device=device), torch.tensor(rois).to(device)),
dim=1)
pool = self.pooling(h, func_roi)
pool = self.conv_pool(pool)
else:
# psroi max align
pool = self.pooling(h, torch.tensor(rois).to(device))
pool = self.flatten(pool)
# fc
fc1 = F.relu(self.fc1(pool))
roi_cls_locs = self.cls_loc(fc1)
roi_scores = self.score(fc1)
return roi_cls_locs, roi_scores
class CNN(Module):
def __init__(self):
super(CNN, self).__init__()
print("[i] Creating cnn layers")
# Input: 4x84x84 tensor
self.layer1 = Conv2d(4, 16, kernel_size=8, stride=4, padding=0)
# Output layer1: 16x20x20 tensor
self.layer3 = Conv2d(16, 32, kernel_size=4, stride=2, padding=0)
# Output layer3: 32x9x9 tensor
self.layer5 = Linear(9 * 9 * 32, 256)
# Output layer5: 256
self.out_layer = Linear(256, 3)
# Output out_layer: 3
print("[i] Initializing layers weights with nn.init.kaiming_normal_")
self.layer1.apply(self.init_weights)
self.layer3.apply(self.init_weights)
self.layer5.apply(self.init_weights)
self.out_layer.apply(self.init_weights)
def init_weights(self, tensor):
init.kaiming_normal_(tensor.weight)
def forward(self, x):
# Converting observation to tensor
x = from_numpy(x).float().to('cuda')
# INPUT TENSORS MUST BE PASSED AS DEPTHxHEIGHTxWIDTH
# Input size: 4x84x84
#print("Input to the cnn: ", x.size())
x = x.unsqueeze(0).to('cuda') # Adding the batch dimension: 1x4x84x84
x = x.permute(0, 3, 2, 1).to('cuda')
#print("Input with batch_dimension ", x.size())
x = x / 255 # Normalizing input
x = F.relu(self.layer1(x)).to(
'cuda') # Conv+ReLU. Input: 1x4x84x84. Output: 1x16x20x20
#print("Conv1 + Relu output size: ", x.size())
x = F.relu(self.layer3(x)).to(
'cuda') # Conv+ReLU Input: 1x16x20x20. Output: 1x32x9x9
#print("Conv2 + Relu output size: ", x.size())
x = x.view(-1, 9 * 9 * 32).to(
'cuda') # Flattening Input: 1x32x9x9. Output: 1x9*9*32
#print("Flattening output size: ", x.size())
x = x.squeeze(0).to('cuda') # Input: 1x32x9x9. Output: 9*9*32
#print("Flattening output size without batch_dimension: ", x.size())
x = F.relu(self.layer5(x)).to(
'cuda') # FC+ReLU Input: 9*9*32. Output: 256
#print("Full connected + Relu output size: ", x.size())
x = self.out_layer(x).to(
'cuda') # Full connected. Input: 256. Output: 4
#print("Out_layer output size: ", x.size())
return x
class LightHeadRCNNResNet101_Head(Module):
def __init__(self,
n_class=81,
roi_size=7,
out_channels=10,
spatial_scale=1 / 16.,
sampling_ratio=2):
super(LightHeadRCNNResNet101_Head, self).__init__()
self.n_class = n_class
self.spatial_scale = spatial_scale
self.roi_size = roi_size
self.out_channels = out_channels
self.sampling_ratio = sampling_ratio
self.c_out = self.roi_size * self.roi_size * self.out_channels
self.global_context_module = GlobalContextModule(
2048, 256, self.c_out, 15)
self.flatten = Flatten()
self.fc1 = Linear(self.c_out, 2048)
self.score = Linear(2048, n_class)
self.cls_loc = Linear(2048, 4 * n_class)
self.apply(lambda x: normal_init(x, 0, 0.01))
self.cls_loc.apply(lambda x: normal_init(x, 0, 0.001))
self.pooling = PSRoIAlign(self.spatial_scale, self.roi_size,
self.sampling_ratio, self.out_channels)
def __call__(self, x, rois):
# global context module
device = x.device
h = self.global_context_module(x)
pool = self.pooling(h, torch.tensor(rois).to(device))
pool = self.flatten(pool)
# fc
fc1 = F.relu(self.fc1(pool))
roi_cls_locs = self.cls_loc(fc1)
roi_scores = self.score(fc1)
return roi_cls_locs, roi_scores
class OptionCriticHead_SharedPreprocessor(Module):
""" Option-critic end output with optional components. Assumes input from shared state preprocessor
Args:
input_size (int): Number of inputs
output_size (int): Number of outputs (per option)
option_size (int): Number of options (O)
intra_option_policy (str): Type of intra-option policy (either discrete or continuous)
intra_option_kwargs (dict, None): Extra args for intra-option policy (e.g., init_log_std and mu_nonlinearity for continuous)
use_interest (bool): If true, apply sigmoid interest function to policy over options (pi_omega)
use_diversity (bool): If true, output a learned per-option Q(s,o) entropy
use_attention (bool): If true, apply a per-option soft attention mechanism to the incoming state. NOT CURRENTLY IMPLEMENTED
baselines_init (bool): If true, use orthogonal initialization
NORM_EPS (float): Normalization constant added to pi_omega normalization to prevent instability.
Outputs:
pi(s): Per-option intra-option policy [BxOxA for discrete, BxO (mu, log_std) for continuous]
beta(s): Per-option termination probability [BxO]
q(s): Per-option value [BxO]
pi_omega(s): Softmax policy over options (parameterized by interest function if using) [BxO]
q_entropy(s): Per-option value entropy (or 0 if not learning) [BxO]
"""
def __init__(self,
input_size: int,
output_size: int,
option_size: int,
intra_option_policy: str,
intra_option_kwargs: [dict, None] = None,
use_interest: bool = False,
use_diversity: bool = False,
use_attention: bool = False,
baselines_init: bool = True,
NORM_EPS: float = 1e-6):
super().__init__()
if intra_option_kwargs is None: intra_option_kwargs = {}
self.use_interest = use_interest
self.use_diversity = use_diversity
self.use_attention = use_attention
self.NORM_EPS = NORM_EPS
pi_class = DiscreteIntraOptionPolicy if intra_option_policy == 'discrete' else ContinuousIntraOptionPolicy
self.pi = pi_class(input_size,
option_size,
output_size,
ortho_init=baselines_init,
**intra_option_kwargs)
self.beta = nn.Sequential(Linear(input_size, option_size),
nn.Sigmoid())
self.q = Linear(input_size, option_size)
self.q_ent = Linear(
input_size, option_size) if use_diversity else Dummy(option_size,
out_value=0.)
self.pi_omega = Sequential(Linear(input_size, option_size),
nn.Softmax(-1))
self.interest = Sequential(
Linear(input_size, option_size),
nn.Sigmoid()) if use_interest else Dummy(option_size)
if baselines_init:
init_v, init_pi = O_INIT_VALUES['v'], O_INIT_VALUES['pi']
self.beta.apply(apply_init)
self.q_ent.apply(apply_init)
self.interest.apply(apply_init)
self.pi_omega.apply(partial(apply_init, gain=init_pi))
self.q.apply(partial(apply_init, gain=init_v))
def forward(self, x):
pi_I = self.pi_omega(x) * self.interest(
x) # Interest-parameterized pi_omega
if self.use_interest: # Avoid messing with computations if we don't have to
pi_I.add_(self.NORM_EPS) # Add eps to avoid instability
pi_I.div_(pi_I.sum(
-1, keepdim=True)) # Normalize so probabilities add to 1
return self.pi(x), self.beta(x), self.q(x), pi_I, self.q_ent(x)
class MRCNer(Module):
"""
基于 MRC 的 ner 模型
"""
def __init__(self, bert_dir: str, dropout: float):
super().__init__()
self.bert = BertModel.from_pretrained(bert_dir)
bert_config = self.bert.config
self.start_classifier = Linear(bert_config.hidden_size, 1)
self.end_classifier = Linear(bert_config.hidden_size, 1)
self.match_classifier = MultiNonLinearClassifier(
bert_config.hidden_size * 2, 1, dropout)
self.init_weights = BertInitWeights(bert_config=bert_config)
self.reset_parameters()
def reset_parameters(self):
self.start_classifier.apply(self.init_weights)
self.end_classifier.apply(self.init_weights)
self.match_classifier.apply(self.init_weights)
def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
token_type_ids: torch.Tensor, sequence_mask: torch.Tensor,
metadata: Dict) -> MRCNerOutput:
"""
模型前向计算
:param input_ids:
:param attention_mask:
:param token_type_ids:
:param sequence_mask:
:param metadata:
:return:
"""
bert_output: BaseModelOutputWithPooling = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
return_dict=True)
sequence_output = bert_output["last_hidden_state"]
sequence_length = sequence_output.size(1)
start_logits = self.start_classifier(sequence_output)
# 最后一个维度去掉 (B, seq_len)
start_logits = start_logits.squeeze(-1)
assert len(start_logits.size()) == 2
end_logits = self.end_classifier(sequence_output)
# (B, seq_len)
end_logits = end_logits.squeeze(-1)
assert len(end_logits.size()) == 2
# 将每一个 i 与 j 连接在一起， 所以是 N*N的拼接，使用了 expand, 进行 两个方向的扩展
# 产生一个 match matrix
# 对于每一个 i 都与 j concat 在一起
# [B, seq_len, seq_len, hidden]
start_extend = sequence_output.unsqueeze(2).expand(
-1, -1, sequence_length, -1)
# [B, seq_len, seq_len, hidden]
end_extend = sequence_output.unsqueeze(1).expand(
-1, sequence_length, -1, -1)
# [B, seq_len, seq_len, hidden*2]
match_matrix = torch.cat([start_extend, end_extend], 3)
# (B, seq_len, seq_len)
match_logits = self.match_classifier(match_matrix).squeeze(-1)
assert len(match_logits.size()) == 3
return MRCNerOutput(start_logits=start_logits,
end_logits=end_logits,
match_logits=match_logits,
mask=sequence_mask)
