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)