neulus/cloud-removal
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pixel shuffle

Browse Source
This commit is contained in:
neulus
2025-10-01 16:30:05 +09:00
parent 8966cafb8f
commit 49025c4d87
19 changed files with 162 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

144
src/model/dino.py
View File

@@ -1,5 +1,6 @@
import math
from typing import Optional, Tuple
import numpy as np
import torch
import torch.nn as nn
@@ -7,7 +8,9 @@ import torch.nn.functional as F
from transformers.models.dinov3_vit.configuration_dinov3_vit import DINOv3ViTConfig
def get_patches_center_coordinates(num_patches_h: int, num_patches_w: int, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
def get_patches_center_coordinates(
num_patches_h: int, num_patches_w: int, dtype: torch.dtype, device: torch.device
) -> torch.Tensor:
coords_h = torch.arange(0.5, num_patches_h, dtype=dtype, device=device)
coords_w = torch.arange(0.5, num_patches_w, dtype=dtype, device=device)
coords_h = coords_h / num_patches_h
@@ -18,8 +21,12 @@ def get_patches_center_coordinates(num_patches_h: int, num_patches_w: int, dtype
return coords
def augment_patches_center_coordinates(coords: torch.Tensor, shift: Optional[float] = None,
jitter: Optional[float] = None, rescale: Optional[float] = None) -> torch.Tensor:
def augment_patches_center_coordinates(
coords: torch.Tensor,
shift: Optional[float] = None,
jitter: Optional[float] = None,
rescale: Optional[float] = None,
) -> torch.Tensor:
if shift is not None:
shift_hw = torch.empty((1, 2), device=coords.device, dtype=coords.dtype)
shift_hw = shift_hw.uniform_(-shift, shift)
@@ -46,7 +53,9 @@ def rotate_half(x):
return torch.cat((-x2, x1), dim=-1)
def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
def apply_rotary_pos_emb(
q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
num_tokens = q.shape[-2]
num_patches = sin.shape[-2]
num_prefix_tokens = num_tokens - num_patches
@@ -63,12 +72,16 @@ def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, si
return q, k
def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
def drop_path(
input: torch.Tensor, drop_prob: float = 0.0, training: bool = False
) -> torch.Tensor:
if drop_prob == 0.0 or not training:
return input
keep_prob = 1 - drop_prob
shape = (input.shape[0],) + (1,) * (input.ndim - 1)
random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
random_tensor = keep_prob + torch.rand(
shape, dtype=input.dtype, device=input.device
)
random_tensor.floor_()
output = input.div(keep_prob) * random_tensor
return output
@@ -80,12 +93,19 @@ class DINOv3ViTEmbeddings(nn.Module):
self.config = config
self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
self.register_tokens = nn.Parameter(torch.empty(1, config.num_register_tokens, config.hidden_size))
self.register_tokens = nn.Parameter(
torch.empty(1, config.num_register_tokens, config.hidden_size)
)
self.patch_embeddings = nn.Conv2d(
config.num_channels, config.hidden_size, kernel_size=config.patch_size, stride=config.patch_size
config.num_channels,
config.hidden_size,
kernel_size=config.patch_size,
stride=config.patch_size,
)
def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.Tensor] = None) -> torch.Tensor:
def forward(
self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.Tensor] = None
) -> torch.Tensor:
batch_size = pixel_values.shape[0]
target_dtype = self.patch_embeddings.weight.dtype
@@ -94,7 +114,9 @@ class DINOv3ViTEmbeddings(nn.Module):
if bool_masked_pos is not None:
mask_token = self.mask_token.to(patch_embeddings.dtype)
patch_embeddings = torch.where(bool_masked_pos.unsqueeze(-1), mask_token, patch_embeddings)
patch_embeddings = torch.where(
bool_masked_pos.unsqueeze(-1), mask_token, patch_embeddings
)
cls_token = self.cls_token.expand(batch_size, -1, -1)
register_tokens = self.register_tokens.expand(batch_size, -1, -1)
@@ -112,7 +134,9 @@ class DINOv3ViTRopePositionEmbedding(nn.Module):
self.num_patches_h = config.image_size // config.patch_size
self.num_patches_w = config.image_size // config.patch_size
inv_freq = 1 / self.base ** torch.arange(0, 1, 4 / self.head_dim, dtype=torch.float32)
inv_freq = 1 / self.base ** torch.arange(
0, 1, 4 / self.head_dim, dtype=torch.float32
)
self.register_buffer("inv_freq", inv_freq, persistent=False)
def forward(self, pixel_values: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
@@ -121,7 +145,11 @@ class DINOv3ViTRopePositionEmbedding(nn.Module):
num_patches_w = width // self.config.patch_size
device = pixel_values.device
device_type = device.type if isinstance(device.type, str) and device.type != "mps" else "cpu"
device_type = (
device.type
if isinstance(device.type, str) and device.type != "mps"
else "cpu"
)
with torch.autocast(device_type=device_type, enabled=False):
patch_coords = get_patches_center_coordinates(
@@ -135,7 +163,9 @@ class DINOv3ViTRopePositionEmbedding(nn.Module):
rescale=self.config.pos_embed_rescale,
)
angles = 2 * math.pi * patch_coords[:, :, None] * self.inv_freq[None, None, :] # type: ignore
angles = (
2 * math.pi * patch_coords[:, :, None] * self.inv_freq[None, None, :] # type: ignore
)
angles = angles.flatten(1, 2)
angles = angles.tile(2)
@@ -161,8 +191,12 @@ class DINOv3ViTAttention(nn.Module):
self.v_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.value_bias)
self.o_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=config.proj_bias)
def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
assert position_embeddings is not None
batch_size, patches, _ = hidden_states.size()
@@ -171,18 +205,32 @@ class DINOv3ViTAttention(nn.Module):
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
query_states = query_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
key_states = key_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
value_states = value_states.view(batch_size, patches, self.num_heads, self.head_dim).transpose(1, 2)
query_states = query_states.view(
batch_size, patches, self.num_heads, self.head_dim
).transpose(1, 2)
key_states = key_states.view(
batch_size, patches, self.num_heads, self.head_dim
).transpose(1, 2)
value_states = value_states.view(
batch_size, patches, self.num_heads, self.head_dim
).transpose(1, 2)
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
query_states, key_states = apply_rotary_pos_emb(
query_states, key_states, cos, sin
)
attn_weights = torch.matmul(query_states, key_states.transpose(-1, -2)) * self.scaling
attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_weights = (
torch.matmul(query_states, key_states.transpose(-1, -2)) * self.scaling
)
attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
query_states.dtype
)
if self.training:
attn_weights = F.dropout(attn_weights, p=self.dropout, training=self.training)
attn_weights = F.dropout(
attn_weights, p=self.dropout, training=self.training
)
if attention_mask is not None:
attn_weights = attn_weights * attention_mask
@@ -198,7 +246,9 @@ class DINOv3ViTAttention(nn.Module):
class DINOv3ViTLayerScale(nn.Module):
def __init__(self, config):
super().__init__()
self.lambda1 = nn.Parameter(config.layerscale_value * torch.ones(config.hidden_size))
self.lambda1 = nn.Parameter(
config.layerscale_value * torch.ones(config.hidden_size)
)
def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
return hidden_state * self.lambda1
@@ -219,8 +269,12 @@ class DINOv3ViTMLP(nn.Module):
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
self.up_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=config.mlp_bias
)
self.down_proj = nn.Linear(
self.intermediate_size, self.hidden_size, bias=config.mlp_bias
)
if config.hidden_act == "gelu":
self.act_fn = F.gelu
@@ -241,9 +295,15 @@ class DINOv3ViTGatedMLP(nn.Module):
self.config = config
self.hidden_size = config.hidden_size
self.intermediate_size = config.intermediate_size
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
self.gate_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=config.mlp_bias
)
self.up_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=config.mlp_bias
)
self.down_proj = nn.Linear(
self.intermediate_size, self.hidden_size, bias=config.mlp_bias
)
if config.hidden_act == "gelu":
self.act_fn = F.gelu
@@ -264,7 +324,11 @@ class DINOv3ViTLayer(nn.Module):
self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self.attention = DINOv3ViTAttention(config)
self.layer_scale1 = DINOv3ViTLayerScale(config)
self.drop_path = DINOv3ViTDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
self.drop_path = (
DINOv3ViTDropPath(config.drop_path_rate)
if config.drop_path_rate > 0.0
else nn.Identity()
)
self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
@@ -274,8 +338,14 @@ class DINOv3ViTLayer(nn.Module):
self.mlp = DINOv3ViTMLP(config)
self.layer_scale2 = DINOv3ViTLayerScale(config)
def forward(self, hidden_states: torch.Tensor, *, attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, **kwargs) -> torch.Tensor:
def forward(
self,
hidden_states: torch.Tensor,
*,
attention_mask: Optional[torch.Tensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
**kwargs,
) -> torch.Tensor:
assert position_embeddings is not None
residual = hidden_states
@@ -303,7 +373,9 @@ class DINOv3ViTModel(nn.Module):
self.config = config
self.embeddings = DINOv3ViTEmbeddings(config)
self.rope_embeddings = DINOv3ViTRopePositionEmbedding(config)
self.layers = nn.ModuleList([DINOv3ViTLayer(config) for _ in range(config.num_hidden_layers)])
self.layers = nn.ModuleList(
[DINOv3ViTLayer(config) for _ in range(config.num_hidden_layers)]
)
self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
self._init_weights()
@@ -337,8 +409,12 @@ class DINOv3ViTModel(nn.Module):
elif isinstance(module, DINOv3ViTLayerScale):
module.lambda1.data.fill_(self.config.layerscale_value)
def forward(self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None):
def forward(
self,
pixel_values: torch.Tensor,
bool_masked_pos: Optional[torch.Tensor] = None,
head_mask: Optional[torch.Tensor] = None,
):
pixel_values = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype)
hidden_states = self.embeddings(pixel_values, bool_masked_pos=bool_masked_pos)
position_embeddings = self.rope_embeddings(pixel_values)