def _init_weights(self, m): if isinstance(m, nn.Linear): trunc_normal_(m.weight, mean=0.0, std=0.02, a=-2.0, b=2.0) if isinstance(m, nn.Linear) and m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.LayerNorm): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0)
def __init__( self, dim: int, num_heads: int, window_size: Sequence[int], qkv_bias: bool = False, attn_drop: float = 0.0, proj_drop: float = 0.0, ) -> None: """ Args: dim: number of feature channels. num_heads: number of attention heads. window_size: local window size. qkv_bias: add a learnable bias to query, key, value. attn_drop: attention dropout rate. proj_drop: dropout rate of output. """ super().__init__() self.dim = dim self.window_size = window_size self.num_heads = num_heads head_dim = dim // num_heads self.scale = head_dim**-0.5 mesh_args = torch.meshgrid.__kwdefaults__ if len(self.window_size) == 3: self.relative_position_bias_table = nn.Parameter( torch.zeros( (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1), num_heads, )) coords_d = torch.arange(self.window_size[0]) coords_h = torch.arange(self.window_size[1]) coords_w = torch.arange(self.window_size[2]) if mesh_args is not None: coords = torch.stack( torch.meshgrid(coords_d, coords_h, coords_w, indexing="ij")) else: coords = torch.stack( torch.meshgrid(coords_d, coords_h, coords_w)) coords_flatten = torch.flatten(coords, 1) relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] relative_coords = relative_coords.permute(1, 2, 0).contiguous() relative_coords[:, :, 0] += self.window_size[0] - 1 relative_coords[:, :, 1] += self.window_size[1] - 1 relative_coords[:, :, 2] += self.window_size[2] - 1 relative_coords[:, :, 0] *= (2 * self.window_size[1] - 1) * (2 * self.window_size[2] - 1) relative_coords[:, :, 1] *= 2 * self.window_size[2] - 1 elif len(self.window_size) == 2: self.relative_position_bias_table = nn.Parameter( torch.zeros( (2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) coords_h = torch.arange(self.window_size[0]) coords_w = torch.arange(self.window_size[1]) if mesh_args is not None: coords = torch.stack( torch.meshgrid(coords_h, coords_w, indexing="ij")) else: coords = torch.stack(torch.meshgrid(coords_h, coords_w)) coords_flatten = torch.flatten(coords, 1) relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] relative_coords = relative_coords.permute(1, 2, 0).contiguous() relative_coords[:, :, 0] += self.window_size[0] - 1 relative_coords[:, :, 1] += self.window_size[1] - 1 relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1 relative_position_index = relative_coords.sum(-1) self.register_buffer("relative_position_index", relative_position_index) self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) self.attn_drop = nn.Dropout(attn_drop) self.proj = nn.Linear(dim, dim) self.proj_drop = nn.Dropout(proj_drop) trunc_normal_(self.relative_position_bias_table, std=0.02) self.softmax = nn.Softmax(dim=-1)
def __init__( self, in_channels: int, img_size: Union[Sequence[int], int], patch_size: Union[Sequence[int], int], hidden_size: int, num_heads: int, pos_embed: str, dropout_rate: float = 0.0, spatial_dims: int = 3, ) -> None: """ Args: in_channels: dimension of input channels. img_size: dimension of input image. patch_size: dimension of patch size. hidden_size: dimension of hidden layer. num_heads: number of attention heads. pos_embed: position embedding layer type. dropout_rate: faction of the input units to drop. spatial_dims: number of spatial dimensions. """ super().__init__() if not (0 <= dropout_rate <= 1): raise ValueError("dropout_rate should be between 0 and 1.") if hidden_size % num_heads != 0: raise ValueError("hidden size should be divisible by num_heads.") self.pos_embed = look_up_option(pos_embed, SUPPORTED_EMBEDDING_TYPES) img_size = ensure_tuple_rep(img_size, spatial_dims) patch_size = ensure_tuple_rep(patch_size, spatial_dims) for m, p in zip(img_size, patch_size): if m < p: raise ValueError("patch_size should be smaller than img_size.") if self.pos_embed == "perceptron" and m % p != 0: raise ValueError( "patch_size should be divisible by img_size for perceptron." ) self.n_patches = np.prod( [im_d // p_d for im_d, p_d in zip(img_size, patch_size)]) self.patch_dim = int(in_channels * np.prod(patch_size)) self.patch_embeddings: nn.Module if self.pos_embed == "conv": self.patch_embeddings = Conv[Conv.CONV, spatial_dims]( in_channels=in_channels, out_channels=hidden_size, kernel_size=patch_size, stride=patch_size) elif self.pos_embed == "perceptron": # for 3d: "b c (h p1) (w p2) (d p3)-> b (h w d) (p1 p2 p3 c)" chars = (("h", "p1"), ("w", "p2"), ("d", "p3"))[:spatial_dims] from_chars = "b c " + " ".join(f"({k} {v})" for k, v in chars) to_chars = f"b ({' '.join([c[0] for c in chars])}) ({' '.join([c[1] for c in chars])} c)" axes_len = {f"p{i+1}": p for i, p in enumerate(patch_size)} self.patch_embeddings = nn.Sequential( Rearrange(f"{from_chars} -> {to_chars}", **axes_len), nn.Linear(self.patch_dim, hidden_size)) self.position_embeddings = nn.Parameter( torch.zeros(1, self.n_patches, hidden_size)) self.dropout = nn.Dropout(dropout_rate) trunc_normal_(self.position_embeddings, mean=0.0, std=0.02, a=-2.0, b=2.0) self.apply(self._init_weights)
def test_ill_arg(self, input_param, input_shape): with self.assertRaises(ValueError): im = torch.rand(input_shape) trunc_normal_(im, **input_param)
def test_shape(self, input_param, input_shape): im = torch.rand(input_shape) trunc_normal_(im, **input_param) self.assertEqual(im.shape, input_shape)