Spaces:
Sleeping
Sleeping
| # Copyright Generate Biomedicines, Inc. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| import torch | |
| import torch.nn as nn | |
| """ | |
| 对3d数据进行归一化 | |
| """ | |
| class MaskedBatchNorm1d(nn.Module): | |
| """A masked version of nn.BatchNorm1d. Only tested for 3D inputs. | |
| Args: | |
| num_features (int): :math:`C` from an expected input of size | |
| :math:`(N, C, L)` | |
| eps (float): a value added to the denominator for numerical stability. | |
| Default: 1e-5 | |
| momentum (float): the value used for the running_mean and running_var | |
| computation. Can be set to ``None`` for cumulative moving average | |
| (i.e. simple average). Default: 0.1 | |
| affine (bool): a boolean value that when set to ``True``, this module has | |
| learnable affine parameters. Default: ``True`` | |
| track_running_stats (bool) : a boolean value that when set to ``True``, this | |
| module tracks the running mean and variance, and when set to ``False``, | |
| this module does not track such statistics and always uses batch | |
| statistics in both training and eval modes. Default: ``True`` | |
| Inputs: | |
| x (torch.tensor): of size (batch_size, num_features, sequence_length) | |
| input_mask (torch.tensor or None) : (optional) of dtype (byte) or (bool) of shape (batch_size, 1, sequence_length) zeroes (or False) indicate positions that cannot contribute to computation | |
| Outputs: | |
| y (torch.tensor): of size (batch_size, num_features, sequence_length) | |
| """ | |
| def __init__( | |
| self, | |
| num_features, | |
| eps=1e-5, | |
| momentum=0.1, | |
| affine=True, | |
| track_running_stats=True, | |
| ): | |
| super(MaskedBatchNorm1d, self).__init__() | |
| self.num_features = num_features | |
| self.eps = eps | |
| self.momentum = momentum | |
| self.affine = affine | |
| if affine: | |
| self.weight = nn.Parameter(torch.Tensor(num_features, 1)) | |
| self.bias = nn.Parameter(torch.Tensor(num_features, 1)) | |
| else: | |
| self.register_parameter("weight", None) | |
| self.register_parameter("bias", None) | |
| self.track_running_stats = track_running_stats | |
| if self.track_running_stats: | |
| self.register_buffer("running_mean", torch.zeros(num_features, 1)) | |
| self.register_buffer("running_var", torch.ones(num_features, 1)) | |
| self.register_buffer( | |
| "num_batches_tracked", torch.tensor(0, dtype=torch.long) | |
| ) | |
| else: | |
| self.register_parameter("running_mean", None) | |
| self.register_parameter("running_var", None) | |
| self.register_parameter("num_batches_tracked", None) | |
| self.reset_parameters() | |
| def reset_running_stats(self): | |
| if self.track_running_stats: | |
| self.running_mean.zero_() | |
| self.running_var.fill_(1) | |
| self.num_batches_tracked.zero_() | |
| def reset_parameters(self): | |
| self.reset_running_stats() | |
| if self.affine: | |
| nn.init.ones_(self.weight) | |
| nn.init.zeros_(self.bias) | |
| def forward(self, input, input_mask=None): | |
| # Calculate the masked mean and variance | |
| B, C, L = input.shape | |
| if input_mask is not None and input_mask.shape != (B, 1, L): | |
| raise ValueError("Mask should have shape (B, 1, L).") | |
| if C != self.num_features: | |
| raise ValueError( | |
| "Expected %d channels but input has %d channels" | |
| % (self.num_features, C) | |
| ) | |
| if input_mask is not None: | |
| masked = input * input_mask | |
| n = input_mask.sum() | |
| else: | |
| masked = input | |
| n = B * L | |
| # Sum | |
| masked_sum = masked.sum(dim=0, keepdim=True).sum(dim=2, keepdim=True) | |
| # Divide by sum of mask | |
| current_mean = masked_sum / n | |
| current_var = ((masked - current_mean) ** 2).sum(dim=0, keepdim=True).sum( | |
| dim=2, keepdim=True | |
| ) / n | |
| # Update running stats | |
| if self.track_running_stats and self.training: | |
| if self.num_batches_tracked == 0: | |
| self.running_mean = current_mean | |
| self.running_var = current_var | |
| else: | |
| self.running_mean = ( | |
| 1 - self.momentum | |
| ) * self.running_mean + self.momentum * current_mean | |
| self.running_var = ( | |
| 1 - self.momentum | |
| ) * self.running_var + self.momentum * current_var | |
| self.num_batches_tracked += 1 | |
| # Norm the input | |
| if self.track_running_stats and not self.training: | |
| normed = (masked - self.running_mean) / ( | |
| torch.sqrt(self.running_var + self.eps) | |
| ) | |
| else: | |
| normed = (masked - current_mean) / (torch.sqrt(current_var + self.eps)) | |
| # Apply affine parameters | |
| if self.affine: | |
| normed = normed * self.weight + self.bias | |
| return normed | |
| class MaskedBatchNorm2d(nn.Module): | |
| """A masked version of nn.BatchNorm1d. Only tested for 3D inputs. | |
| Args: | |
| num_features (int): :math:`C` from an expected input of size | |
| :math:`(N, C, L)` | |
| eps (float): a value added to the denominator for numerical stability. | |
| Default: 1e-5 | |
| momentum (float): the value used for the running_mean and running_var | |
| computation. Can be set to ``None`` for cumulative moving average | |
| (i.e. simple average). Default: 0.1 | |
| affine (bool): a boolean value that when set to ``True``, this module has | |
| learnable affine parameters. Default: ``True`` | |
| track_running_stats (bool) : a boolean value that when set to ``True``, this | |
| module tracks the running mean and variance, and when set to ``False``, | |
| this module does not track such statistics and always uses batch | |
| statistics in both training and eval modes. Default: ``True`` | |
| Inputs: | |
| x (torch.tensor): of size (batch_size, num_features, sequence_length) | |
| input_mask (torch.tensor or None) : (optional) of dtype (byte) or (bool) of shape (batch_size, 1, sequence_length) zeroes (or False) indicate positions that cannot contribute to computation | |
| Outputs: | |
| y (torch.tensor): of size (batch_size, num_features, sequence_length) | |
| """ | |
| def __init__( | |
| self, | |
| num_features, | |
| eps=1e-5, | |
| momentum=0.1, | |
| affine=True, | |
| track_running_stats=True, | |
| ): | |
| super().__init__() | |
| self.num_features = num_features | |
| self.eps = eps | |
| self.momentum = momentum | |
| self.affine = affine | |
| if affine: | |
| self.weight = nn.Parameter(torch.ones(num_features,)) | |
| self.bias = nn.Parameter(torch.zeros(num_features,)) | |
| else: | |
| self.register_parameter("weight", None) | |
| self.register_parameter("bias", None) | |
| self.track_running_stats = track_running_stats | |
| if self.track_running_stats: | |
| self.register_buffer("running_mean", torch.zeros(1, 1, 1, num_features)) | |
| self.register_buffer("running_var", torch.ones(1, 1, 1, num_features)) | |
| self.register_buffer( | |
| "num_batches_tracked", torch.tensor(0, dtype=torch.long) | |
| ) | |
| else: | |
| self.register_parameter("running_mean", None) | |
| self.register_parameter("running_var", None) | |
| self.register_parameter("num_batches_tracked", None) | |
| self.reset_parameters() | |
| def reset_running_stats(self): | |
| if self.track_running_stats: | |
| self.running_mean.zero_() | |
| self.running_var.fill_(1) | |
| self.num_batches_tracked.zero_() | |
| def reset_parameters(self): | |
| self.reset_running_stats() | |
| if self.affine: | |
| nn.init.ones_(self.weight) | |
| nn.init.zeros_(self.bias) | |
| def forward(self, input, mask=None): | |
| # Calculate the masked mean and variance | |
| B, L, L, C = input.size() | |
| if mask is not None: | |
| if mask.dim() != 4: | |
| raise ValueError( | |
| f"Input mask must have four dimensions, but has {mask.dim()}" | |
| ) | |
| b, l, l, d = mask.size() | |
| if (b != B) or (l != L): | |
| raise ValueError( | |
| f"Input mask must have shape {(B, L, L, 1)} or {(B, L, L, C)} to match input." | |
| ) | |
| if d == 1: | |
| mask = mask.expand(input.size()) | |
| if C != self.num_features: | |
| raise ValueError( | |
| "Expected %d channels but input has %d channels" | |
| % (self.num_features, C) | |
| ) | |
| if mask is None: | |
| mask = input.new_ones(input.size()) | |
| masked = input * mask | |
| n = mask.sum(dim=(0, 1, 2), keepdim=True) | |
| masked_sum = (masked).sum(dim=(0, 1, 2), keepdim=True) | |
| current_mean = masked_sum / n | |
| current_var = (mask * (masked - current_mean).pow(2)).sum( | |
| dim=(0, 1, 2), keepdim=True | |
| ) / n | |
| # Update running stats | |
| with torch.no_grad(): | |
| if self.track_running_stats and self.training: | |
| if self.num_batches_tracked == 0: | |
| self.running_mean = current_mean.detach() | |
| self.running_var = current_var.detach() | |
| else: | |
| self.running_mean = ( | |
| 1 - self.momentum | |
| ) * self.running_mean + self.momentum * current_mean.detach() | |
| self.running_var = ( | |
| 1 - self.momentum | |
| ) * self.running_var + self.momentum * current_var.detach() | |
| self.num_batches_tracked += 1 | |
| # Norm the input | |
| if self.track_running_stats and not self.training: | |
| normed = (masked - self.running_mean) / ( | |
| torch.sqrt(self.running_var + self.eps) | |
| ) | |
| else: | |
| normed = (masked - current_mean) / (torch.sqrt(current_var + self.eps)) | |
| # Apply affine parameters | |
| if self.affine: | |
| normed = normed * self.weight + self.bias | |
| normed = normed * mask | |
| return normed | |
| class NormedReductionLayer(nn.Module): | |
| """A ReductionLayer with LayerNorms after the hidden layers.""" | |
| def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.0): | |
| super().__init__() | |
| self.d1 = nn.Dropout(p=dropout) | |
| self.d2 = nn.Dropout(p=dropout) | |
| self.hidden = nn.Linear(input_dim, hidden_dim) | |
| self.relu = nn.ReLU() | |
| self.output = nn.Linear(hidden_dim, output_dim) | |
| self.norm1 = nn.LayerNorm(input_dim) | |
| self.norm2 = nn.LayerNorm(hidden_dim) | |
| def reduce(self, x, mask): | |
| masked_x = x * mask | |
| mean_x = masked_x.sum(dim=1) / torch.sum(mask, dim=1) | |
| return mean_x | |
| def forward(self, x, mask): | |
| reduced_x = self.norm1(self.reduce(x, mask)) | |
| h = self.norm2(self.hidden(reduced_x)) | |
| return self.output(self.relu(h)) | |