jhqxxx
1381 lines
44 KiB
Rust
1381 lines
44 KiB
Rust
use anyhow::{Result, anyhow};
|
|
use candle_core::{D, DType, Device, IndexOp, Shape, Tensor};
|
|
use candle_nn::{
|
|
Activation, BatchNorm, Conv2d, Init, LayerNorm, Linear, Module, ModuleT, VarBuilder, linear,
|
|
linear_b, linear_no_bias, ops::sigmoid,
|
|
};
|
|
|
|
use crate::{
|
|
models::common::{
|
|
Conv2dWithBN, TwoLinearMLP, deform_conv2d_kernel, get_batch_norm, get_conv2d,
|
|
get_layer_norm,
|
|
},
|
|
utils::tensor_utils::{
|
|
get_equal_mask, index_select_2d, interpolate_bilinear, split_tensor_with_size,
|
|
},
|
|
};
|
|
|
|
struct PatchEmbed {
|
|
proj: Conv2d,
|
|
norm: Option<LayerNorm>,
|
|
patch_size: usize,
|
|
embed_dim: usize,
|
|
}
|
|
|
|
impl PatchEmbed {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
in_chans: usize,
|
|
embed_dim: usize,
|
|
patch_size: usize,
|
|
patch_norm: bool,
|
|
) -> Result<Self> {
|
|
let proj = get_conv2d(
|
|
vb.pp("proj"),
|
|
in_chans,
|
|
embed_dim,
|
|
patch_size,
|
|
0,
|
|
patch_size,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let norm = if patch_norm {
|
|
Some(get_layer_norm(vb.pp("norm"), 1e-5, embed_dim)?)
|
|
} else {
|
|
None
|
|
};
|
|
Ok(Self {
|
|
patch_size,
|
|
proj,
|
|
norm,
|
|
embed_dim,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let (_, _, h, w) = xs.dims4()?;
|
|
let mut xs = xs.clone();
|
|
if w % self.patch_size != 0 {
|
|
xs = xs.pad_with_zeros(3, 0, self.patch_size - w % self.patch_size)?;
|
|
}
|
|
if h % self.patch_size != 0 {
|
|
xs = xs.pad_with_zeros(2, 0, self.patch_size - h % self.patch_size)?;
|
|
}
|
|
xs = self.proj.forward(&xs)?;
|
|
if self.norm.is_some() {
|
|
let (_, _, ph, pw) = xs.dims4()?;
|
|
xs = xs.flatten_from(2)?.transpose(1, 2)?;
|
|
xs = self.norm.as_ref().unwrap().forward(&xs)?;
|
|
xs = xs.transpose(1, 2)?.reshape(((), self.embed_dim, ph, pw))?;
|
|
}
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
pub struct WindowAttention {
|
|
num_heads: usize,
|
|
relative_position_bias: Tensor,
|
|
qkv: Linear,
|
|
proj: Linear,
|
|
scaling: f64,
|
|
}
|
|
|
|
impl WindowAttention {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
dim: usize,
|
|
num_heads: usize,
|
|
qkv_bias: bool,
|
|
window_size: (usize, usize),
|
|
) -> Result<Self> {
|
|
let head_dim = dim / num_heads;
|
|
let scaling = 1.0 / (head_dim as f64).sqrt();
|
|
let qkv = linear_b(dim, dim * 3, qkv_bias, vb.pp("qkv"))?;
|
|
|
|
let proj = linear(dim, dim, vb.pp("proj"))?;
|
|
let relative_position_bias_table = vb.get_with_hints(
|
|
((2 * window_size.0 - 1) * (2 * window_size.1 - 1), num_heads),
|
|
"relative_position_bias_table",
|
|
Init::Const(0.),
|
|
)?; //2*Wh-1 * 2*Ww-1, nH
|
|
|
|
let coords_h = Tensor::arange(0f32, window_size.0 as f32, vb.device())?
|
|
.unsqueeze(1)?
|
|
.broadcast_as(window_size)?;
|
|
let coords_w = Tensor::arange(0f32, window_size.1 as f32, vb.device())?
|
|
.unsqueeze(0)?
|
|
.broadcast_as(window_size)?;
|
|
|
|
let coords = Tensor::stack(&[coords_h, coords_w], 0)?.flatten_from(1)?; // (2, wh, ww)
|
|
let coords1 = coords.unsqueeze(2)?;
|
|
let coords2 = coords.unsqueeze(1)?;
|
|
let relative_coords = coords1
|
|
.broadcast_sub(&coords2)?
|
|
.permute((1, 2, 0))?
|
|
.contiguous()?; // (wh*ww, wh*ww, 2)
|
|
let relative_coords_0 = relative_coords
|
|
.i((.., .., 0))?
|
|
.affine(1.0, window_size.0 as f64 - 1.0)?;
|
|
let relative_coords_1 = relative_coords
|
|
.i((.., .., 1))?
|
|
.affine(1.0, window_size.1 as f64 - 1.0)?;
|
|
let relative_coords_0 = relative_coords_0.affine(2.0 * window_size.1 as f64 - 1.0, 0.0)?;
|
|
let relative_position_index = relative_coords_0
|
|
.add(&relative_coords_1)?
|
|
.to_dtype(candle_core::DType::U32)?;
|
|
let relative_position_bias =
|
|
index_select_2d(&relative_position_bias_table, &relative_position_index)?;
|
|
Ok(Self {
|
|
num_heads,
|
|
relative_position_bias,
|
|
qkv,
|
|
proj,
|
|
scaling,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor, attn_mask: Option<&Tensor>) -> Result<Tensor> {
|
|
let (b, seq_len, _) = xs.dims3()?;
|
|
// (3, B, n_head, h*w, head_dim)
|
|
let qkv = self
|
|
.qkv
|
|
.forward(xs)?
|
|
.reshape((b, seq_len, 3, self.num_heads, ()))?
|
|
.permute((2, 0, 3, 1, 4))?
|
|
.contiguous()?;
|
|
let query_states = qkv.i(0)?.contiguous()?;
|
|
let key_states = qkv.i(1)?.contiguous()?;
|
|
let value_states = qkv.i(2)?.contiguous()?;
|
|
let attn_bias = self
|
|
.relative_position_bias
|
|
.permute((2, 0, 1))?
|
|
.contiguous()?
|
|
.unsqueeze(0)?;
|
|
let query_states = (query_states * self.scaling)?;
|
|
|
|
let attn_weights = query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?;
|
|
let attn_weights = attn_weights.broadcast_add(&attn_bias)?;
|
|
let attn_weights = match attn_mask {
|
|
None => attn_weights,
|
|
Some(mask) => {
|
|
let nw: usize = mask.dim(0)?;
|
|
let attn_weights = attn_weights
|
|
.reshape((b / nw, nw, self.num_heads, seq_len, seq_len))?
|
|
.broadcast_add(
|
|
&mask
|
|
.unsqueeze(1)?
|
|
.unsqueeze(0)?
|
|
.to_dtype(attn_weights.dtype())?,
|
|
)?;
|
|
|
|
attn_weights.reshape(((), self.num_heads, seq_len, seq_len))?
|
|
}
|
|
};
|
|
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
|
let attn_output = attn_weights.matmul(&value_states)?;
|
|
|
|
//(b, n_head, seq_len, dim) -> (b, seq_len, n_head, dim)
|
|
let xs = attn_output.transpose(1, 2)?.contiguous()?;
|
|
// (b, h*w, n_head, dim)
|
|
let xs = xs.reshape((b, seq_len, ()))?;
|
|
let xs = self.proj.forward(&xs)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
fn window_partition(x: &Tensor, window_size: usize) -> Result<Tensor> {
|
|
let (b, h, w, c) = x.dims4()?;
|
|
|
|
let x = x.reshape((
|
|
b,
|
|
h / window_size,
|
|
window_size,
|
|
w / window_size,
|
|
window_size,
|
|
c,
|
|
))?;
|
|
let windows =
|
|
x.permute((0, 1, 3, 2, 4, 5))?
|
|
.contiguous()?
|
|
.reshape(((), window_size, window_size, c))?;
|
|
Ok(windows)
|
|
}
|
|
|
|
fn window_reverse(windows: &Tensor, window_size: usize, pad_hw: (usize, usize)) -> Result<Tensor> {
|
|
let (hp, wp) = pad_hw;
|
|
let b = windows.dim(0)? / (hp * wp / window_size / window_size);
|
|
let last_dim = windows.dim(D::Minus1)?;
|
|
let x = windows.reshape(&[
|
|
b,
|
|
hp / window_size,
|
|
wp / window_size,
|
|
window_size,
|
|
window_size,
|
|
last_dim,
|
|
])?;
|
|
let x = x
|
|
.permute((0, 1, 3, 2, 4, 5))?
|
|
.contiguous()?
|
|
.reshape((b, hp, wp, ()))?;
|
|
Ok(x)
|
|
}
|
|
|
|
struct SwinTransformerBlock {
|
|
norm1: LayerNorm,
|
|
attn: WindowAttention,
|
|
norm2: LayerNorm,
|
|
mlp: TwoLinearMLP,
|
|
window_size: usize,
|
|
shift_size: usize,
|
|
}
|
|
|
|
impl SwinTransformerBlock {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
dim: usize,
|
|
num_heads: usize,
|
|
mlp_ratio: f32,
|
|
qkv_bias: bool,
|
|
act: Activation,
|
|
window_size: usize,
|
|
shift_size: usize,
|
|
) -> Result<Self> {
|
|
let norm1 = get_layer_norm(vb.pp("norm1"), 1e-5, dim)?;
|
|
|
|
let attn = WindowAttention::new(
|
|
vb.pp("attn"),
|
|
dim,
|
|
num_heads,
|
|
qkv_bias,
|
|
(window_size, window_size),
|
|
)?;
|
|
let norm2 = get_layer_norm(vb.pp("norm2"), 1e-5, dim)?;
|
|
let mlp_dim = (dim as f32 * mlp_ratio) as usize;
|
|
let mlp = TwoLinearMLP::new(vb.pp("mlp"), dim, mlp_dim, dim, act, true, "fc1", "fc2")?;
|
|
Ok(Self {
|
|
norm1,
|
|
attn,
|
|
norm2,
|
|
mlp,
|
|
window_size,
|
|
shift_size,
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&self,
|
|
xs: &Tensor,
|
|
mask_matrix: Option<&Tensor>,
|
|
h: usize,
|
|
w: usize,
|
|
) -> Result<Tensor> {
|
|
let (b, seq_len, c) = xs.dims3()?;
|
|
assert_eq!(
|
|
seq_len,
|
|
h * w,
|
|
"swin transformer block sq_len not equal to h*w"
|
|
);
|
|
let shortcut = xs.clone();
|
|
let xs = self.norm1.forward(xs)?;
|
|
let xs = xs.reshape((b, h, w, c))?;
|
|
let pad_h = (self.window_size - h % self.window_size) % self.window_size;
|
|
let pad_w = (self.window_size - w % self.window_size) % self.window_size;
|
|
let xs = xs.pad_with_zeros(1, 0, pad_h)?;
|
|
let xs = xs.pad_with_zeros(2, 0, pad_w)?;
|
|
let (_, hp, wp, _) = xs.dims4()?;
|
|
|
|
let (shifted_x, attn_mask) = if self.shift_size > 0 {
|
|
(
|
|
xs.roll(-(self.shift_size as i32), 1)?
|
|
.roll(-(self.shift_size as i32), 2)?,
|
|
mask_matrix,
|
|
)
|
|
} else {
|
|
(xs, None)
|
|
};
|
|
let xs = window_partition(&shifted_x, self.window_size)?;
|
|
let xs = xs.reshape(((), self.window_size * self.window_size, c))?;
|
|
let xs = self.attn.forward(&xs, attn_mask)?;
|
|
let xs = window_reverse(&xs, self.window_size, (hp, wp))?;
|
|
let mut xs = if self.shift_size > 0 {
|
|
xs.roll(self.shift_size as i32, 1)?
|
|
.roll(self.shift_size as i32, 2)?
|
|
} else {
|
|
xs
|
|
};
|
|
if pad_h > 0 || pad_w > 0 {
|
|
xs = xs.i((.., 0..h, 0..w, ..))?.contiguous()?;
|
|
}
|
|
let xs = xs.reshape((b, h * w, c))?;
|
|
let x = shortcut.add(&xs)?;
|
|
let x = x.add(&self.mlp.forward(&self.norm2.forward(&x)?)?)?;
|
|
Ok(x)
|
|
}
|
|
}
|
|
|
|
struct PatchMerging {
|
|
reduction: Linear,
|
|
norm: LayerNorm,
|
|
}
|
|
|
|
impl PatchMerging {
|
|
pub fn new(vb: VarBuilder, dim: usize) -> Result<Self> {
|
|
let reduction = linear_no_bias(4 * dim, 2 * dim, vb.pp("reduction"))?;
|
|
let norm = get_layer_norm(vb.pp("norm"), 1e-5, 4 * dim)?;
|
|
Ok(Self { reduction, norm })
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor, h: usize, w: usize) -> Result<Tensor> {
|
|
let (b, l, c) = xs.dims3()?;
|
|
assert_eq!(l, h * w, "input feature has wrong size");
|
|
let mut xs = xs.reshape((b, h, w, c))?;
|
|
let pad_input = (h % 2 == 1) || (w % 2 == 1);
|
|
if pad_input {
|
|
xs = xs
|
|
.pad_with_zeros(2, 0, w % 2)?
|
|
.pad_with_zeros(1, 0, h % 2)?;
|
|
}
|
|
let shape = Shape::from_dims(&[b, h / 2, 2, w / 2, 2, c]);
|
|
let xs = xs.reshape(shape)?;
|
|
let x0 = xs.i((.., .., 0, .., 0, ..))?;
|
|
let x1 = xs.i((.., .., 1, .., 0, ..))?;
|
|
let x2 = xs.i((.., .., 0, .., 1, ..))?;
|
|
let x3 = xs.i((.., .., 1, .., 1, ..))?;
|
|
let xs = Tensor::cat(&[x0, x1, x2, x3], D::Minus1)?;
|
|
let xs = xs.reshape((b, (), 4 * c))?;
|
|
let xs = self.norm.forward(&xs)?;
|
|
let xs = self.reduction.forward(&xs)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
struct BasicLayer {
|
|
window_size: usize,
|
|
shift_size: usize,
|
|
blocks: Vec<SwinTransformerBlock>,
|
|
downsample: Option<PatchMerging>,
|
|
}
|
|
|
|
impl BasicLayer {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
dim: usize,
|
|
depth: usize,
|
|
num_heads: usize,
|
|
window_size: usize,
|
|
mlp_ratio: f32,
|
|
qkv_bias: bool,
|
|
downsample: bool,
|
|
) -> Result<Self> {
|
|
let shift_size = window_size / 2;
|
|
let mut blocks = vec![];
|
|
let vb_blocks = vb.pp("blocks");
|
|
for i in 0..depth {
|
|
let block_shift_size = if i % 2 == 0 { 0usize } else { shift_size };
|
|
|
|
let block = SwinTransformerBlock::new(
|
|
vb_blocks.pp(i),
|
|
dim,
|
|
num_heads,
|
|
mlp_ratio,
|
|
qkv_bias,
|
|
Activation::Gelu,
|
|
window_size,
|
|
block_shift_size,
|
|
)?;
|
|
blocks.push(block);
|
|
}
|
|
let downsample = if downsample {
|
|
Some(PatchMerging::new(vb.pp("downsample"), dim)?)
|
|
} else {
|
|
None
|
|
};
|
|
Ok(Self {
|
|
window_size,
|
|
shift_size,
|
|
blocks,
|
|
downsample,
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&self,
|
|
xs: &Tensor,
|
|
h: usize,
|
|
w: usize,
|
|
) -> Result<(Tensor, usize, usize, Tensor, usize, usize)> {
|
|
let hp = (h as f32 / self.window_size as f32).ceil() as usize * self.window_size;
|
|
let wp = (w as f32 / self.window_size as f32).ceil() as usize * self.window_size;
|
|
let mut img_mask = Tensor::zeros((1, hp, wp, 1), xs.dtype(), xs.device())?;
|
|
let h_slices = [
|
|
(0usize, hp - self.window_size),
|
|
(hp - self.window_size, hp - self.shift_size),
|
|
(hp - self.shift_size, hp),
|
|
];
|
|
let w_slices = [
|
|
(0usize, wp - self.window_size),
|
|
(wp - self.window_size, wp - self.shift_size),
|
|
(wp - self.shift_size, wp),
|
|
];
|
|
let mut cnt = 0f64;
|
|
for (h_start, h_end) in h_slices {
|
|
for (w_start, w_end) in w_slices {
|
|
let mask_value = Tensor::zeros(
|
|
(1, h_end - h_start, w_end - w_start, 1),
|
|
xs.dtype(),
|
|
xs.device(),
|
|
)?
|
|
.affine(1.0, cnt)?;
|
|
img_mask = img_mask.slice_assign(
|
|
&[(0..1), (h_start..h_end), (w_start..w_end), (0..1)],
|
|
&mask_value,
|
|
)?;
|
|
cnt += 1.0;
|
|
}
|
|
}
|
|
|
|
let mask_windows = window_partition(&img_mask, self.window_size)?;
|
|
let mask_windows = mask_windows.reshape(((), self.window_size * self.window_size))?;
|
|
let attn_mask = mask_windows
|
|
.unsqueeze(1)?
|
|
.broadcast_sub(&mask_windows.unsqueeze(2)?)?;
|
|
let equal_zero_mask = get_equal_mask(&attn_mask, 0)?;
|
|
let attn_mask = equal_zero_mask.where_cond(
|
|
&Tensor::new(0f32, xs.device())?.broadcast_as(equal_zero_mask.shape())?,
|
|
&Tensor::new(-100f32, xs.device())?.broadcast_as(equal_zero_mask.shape())?,
|
|
)?;
|
|
let mut xs = xs.clone();
|
|
for block in &self.blocks {
|
|
xs = block.forward(&xs, Some(&attn_mask), h, w)?;
|
|
}
|
|
let (xs_down, wh, ww) = match self.downsample.as_ref() {
|
|
Some(down) => {
|
|
let xs_down = down.forward(&xs, h, w)?;
|
|
// let wh = (h + 1) / 2;
|
|
// let ww = (w + 1) / 2;
|
|
let wh = h.div_ceil(2);
|
|
let ww = w.div_ceil(2);
|
|
(xs_down, wh, ww)
|
|
}
|
|
None => (xs.clone(), h, w),
|
|
};
|
|
Ok((xs, h, w, xs_down, wh, ww))
|
|
}
|
|
}
|
|
|
|
pub struct SwinTransformer {
|
|
patch_embed: PatchEmbed,
|
|
num_layers: usize,
|
|
// pos_drop: Dropout,
|
|
layers: Vec<BasicLayer>,
|
|
norms: Vec<LayerNorm>,
|
|
out_indices: Vec<usize>,
|
|
num_features: Vec<usize>,
|
|
}
|
|
|
|
impl SwinTransformer {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
patch_size: usize,
|
|
in_channels: usize,
|
|
embed_dim: usize,
|
|
depths: Vec<usize>,
|
|
num_heads: Vec<usize>,
|
|
window_size: usize,
|
|
mlp_ratio: f32,
|
|
qkv_bias: bool,
|
|
patch_norm: bool,
|
|
out_indices: Vec<usize>,
|
|
) -> Result<Self> {
|
|
let patch_embed = PatchEmbed::new(
|
|
vb.pp("patch_embed"),
|
|
in_channels,
|
|
embed_dim,
|
|
patch_size,
|
|
patch_norm,
|
|
)?;
|
|
let num_layers = depths.len();
|
|
let mut layers = vec![];
|
|
let vb_layers = vb.pp("layers");
|
|
let mut num_features = vec![];
|
|
for i in 0..num_layers {
|
|
let downsample = i < num_layers - 1;
|
|
let dim_i = embed_dim * 2usize.pow(i as u32);
|
|
num_features.push(dim_i);
|
|
let layer_i = BasicLayer::new(
|
|
vb_layers.pp(i),
|
|
dim_i,
|
|
depths[i],
|
|
num_heads[i],
|
|
window_size,
|
|
mlp_ratio,
|
|
qkv_bias,
|
|
downsample,
|
|
)?;
|
|
layers.push(layer_i);
|
|
}
|
|
let mut norms = vec![];
|
|
for i in out_indices.clone() {
|
|
let layer_i = get_layer_norm(vb.pp(format!("norm{i}")), 1e-5, num_features[i])?;
|
|
norms.push(layer_i);
|
|
}
|
|
Ok(Self {
|
|
num_layers,
|
|
patch_embed,
|
|
layers,
|
|
norms,
|
|
out_indices,
|
|
num_features,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Vec<Tensor>> {
|
|
let xs = self.patch_embed.forward(xs)?;
|
|
let (_, _, mut wh, mut ww) = xs.dims4()?;
|
|
let mut outs = vec![];
|
|
let mut xs = xs.flatten_from(2)?.transpose(1, 2)?;
|
|
let mut norm_idx = 0;
|
|
for i in 0..self.num_layers {
|
|
let layer = &self.layers[i];
|
|
let (x_out, h, w, xs_, wh_, ww_) = layer.forward(&xs, wh, ww)?;
|
|
xs = xs_.clone();
|
|
wh = wh_;
|
|
ww = ww_;
|
|
if self.out_indices.contains(&i) {
|
|
let norm_layer = &self.norms[norm_idx];
|
|
norm_idx += 1;
|
|
let x_out = norm_layer.forward(&x_out)?;
|
|
let out = x_out
|
|
.reshape(((), h, w, self.num_features[i]))?
|
|
.permute((0, 3, 1, 2))?
|
|
.contiguous()?;
|
|
outs.push(out);
|
|
}
|
|
}
|
|
Ok(outs)
|
|
}
|
|
}
|
|
|
|
#[allow(unused)]
|
|
struct DeformableConv2d {
|
|
offset_conv: Conv2d,
|
|
modulator_conv: Conv2d,
|
|
regular_conv: Conv2d,
|
|
stride: usize,
|
|
padding: usize,
|
|
ks: usize,
|
|
}
|
|
|
|
#[allow(unused)]
|
|
impl DeformableConv2d {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
in_c: usize,
|
|
out_c: usize,
|
|
kernel_size: usize,
|
|
stride: usize,
|
|
padding: usize,
|
|
bias: bool,
|
|
) -> Result<Self> {
|
|
let offset_conv = get_conv2d(
|
|
vb.pp("offset_conv"),
|
|
in_c,
|
|
2 * kernel_size * kernel_size,
|
|
kernel_size,
|
|
padding,
|
|
stride,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
|
|
let modulator_conv = get_conv2d(
|
|
vb.pp("modulator_conv"),
|
|
in_c,
|
|
kernel_size * kernel_size,
|
|
kernel_size,
|
|
padding,
|
|
stride,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
|
|
let regular_conv = get_conv2d(
|
|
vb.pp("regular_conv"),
|
|
in_c,
|
|
out_c,
|
|
kernel_size,
|
|
0,
|
|
kernel_size,
|
|
1,
|
|
1,
|
|
bias,
|
|
)?;
|
|
Ok(Self {
|
|
offset_conv,
|
|
modulator_conv,
|
|
regular_conv,
|
|
stride,
|
|
padding,
|
|
ks: kernel_size,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
|
self.forward_use_kernel(xs)
|
|
// if self.ks > 1 {
|
|
// self.forward_use_kernel(xs)
|
|
// } else {
|
|
// self.forward_use_tensor(xs)
|
|
// }
|
|
}
|
|
pub fn forward_use_kernel(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let offset = self.offset_conv.forward(xs)?; // (b, 2*k*k, out_h, out_w)
|
|
|
|
let modulator = sigmoid(&self.modulator_conv.forward(xs)?)?
|
|
.affine(2.0, 0.0)?
|
|
.contiguous()?;
|
|
let out = deform_conv2d_kernel(
|
|
xs,
|
|
self.regular_conv.weight(),
|
|
self.regular_conv.bias(),
|
|
&offset,
|
|
Some(&modulator),
|
|
self.stride,
|
|
self.padding,
|
|
)?;
|
|
Ok(out)
|
|
}
|
|
|
|
pub fn forward_use_tensor(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let offset = self.offset_conv.forward(xs)?; // (b, 2*k*k, out_h, out_w)
|
|
|
|
let modulator = sigmoid(&self.modulator_conv.forward(xs)?)?
|
|
.affine(2.0, 0.0)?
|
|
.contiguous()?;
|
|
let n = offset.dim(1)? / 2;
|
|
|
|
let xs = if self.padding > 0 {
|
|
xs.pad_with_zeros(2, self.padding, self.padding)?
|
|
.pad_with_zeros(3, self.padding, self.padding)?
|
|
} else {
|
|
xs.clone()
|
|
};
|
|
let offset = if self.ks > 3 {
|
|
offset.to_device(&Device::Cpu)?
|
|
} else {
|
|
offset
|
|
};
|
|
let p = self.get_p(&offset)?;
|
|
// drop(offset);
|
|
// (b, h, w, 2n)
|
|
let p = p.permute((0, 2, 3, 1))?.contiguous()?;
|
|
let q_lt = p.floor()?;
|
|
let q_rb = (&q_lt + 1.0)?;
|
|
let (_, _, in_h, in_w) = xs.dims4()?;
|
|
let in_h = in_h as f64;
|
|
let in_w = in_w as f64;
|
|
|
|
// 分开处理x和y坐标
|
|
let p_x = p.narrow(3, 0, n)?.clamp(0.0, in_h - 1.0)?;
|
|
let p_y = p.narrow(3, n, n)?.clamp(0.0, in_w - 1.0)?;
|
|
// drop(p);
|
|
let q_lt_x = q_lt.narrow(3, 0, n)?.clamp(0.0, in_h - 1.0)?;
|
|
let q_lt_y = q_lt.narrow(3, n, n)?.clamp(0.0, in_w - 1.0)?;
|
|
let q_rb_x = q_rb.narrow(3, 0, n)?.clamp(0.0, in_h - 1.0)?;
|
|
let q_rb_y = q_rb.narrow(3, n, n)?.clamp(0.0, in_w - 1.0)?;
|
|
// drop(q_lt);
|
|
// drop(q_rb);
|
|
// 转换为整数索引
|
|
let q_lt_x_idx = q_lt_x.to_dtype(DType::U32)?;
|
|
let q_lt_y_idx = q_lt_y.to_dtype(DType::U32)?;
|
|
let q_rb_x_idx = q_rb_x.to_dtype(DType::U32)?;
|
|
let q_rb_y_idx = q_rb_y.to_dtype(DType::U32)?;
|
|
|
|
// 计算双线性权重
|
|
let p_sub_lt_x = (&p_x - &q_lt_x)?;
|
|
let one_sub_lt_x = (1.0 - &p_sub_lt_x)?;
|
|
let p_sub_lt_y = (&p_y - &q_lt_y)?;
|
|
let one_sub_lt_y = (1.0 - &p_sub_lt_y)?;
|
|
// drop(q_lt_x);
|
|
// drop(q_lt_y);
|
|
// drop(q_rb_x);
|
|
// drop(q_rb_y);
|
|
let g_lt = (&one_sub_lt_x * &one_sub_lt_y)?;
|
|
let g_rb = (&p_sub_lt_x * &p_sub_lt_y)?;
|
|
let g_lb = (&one_sub_lt_x * &p_sub_lt_y)?;
|
|
let g_rt = (&p_sub_lt_x * &one_sub_lt_y)?;
|
|
// drop(p_sub_lt_x);
|
|
// drop(one_sub_lt_x);
|
|
// drop(p_sub_lt_y);
|
|
// drop(one_sub_lt_y);
|
|
|
|
let xs = if self.ks > 3 {
|
|
xs.to_device(&Device::Cpu)?
|
|
} else {
|
|
xs
|
|
};
|
|
// 采样四个角点的特征
|
|
let x_q_lt = self.get_x_q(&xs, &q_lt_x_idx, &q_lt_y_idx)?;
|
|
let x_q_rb = self.get_x_q(&xs, &q_rb_x_idx, &q_rb_y_idx)?;
|
|
let x_q_lb = self.get_x_q(&xs, &q_lt_x_idx, &q_rb_y_idx)?;
|
|
let x_q_rt = self.get_x_q(&xs, &q_rb_x_idx, &q_lt_y_idx)?;
|
|
// drop(q_lt_x_idx);
|
|
// drop(q_lt_y_idx);
|
|
// drop(q_rb_x_idx);
|
|
// drop(q_rb_y_idx);
|
|
// 双线性插值
|
|
let x_offset = g_lt.unsqueeze(1)?.broadcast_mul(&x_q_lt)?;
|
|
// drop(g_lt);
|
|
// drop(x_q_lt);
|
|
let x_offset = x_offset.add(&g_rb.unsqueeze(1)?.broadcast_mul(&x_q_rb)?)?;
|
|
// drop(g_rb);
|
|
// drop(x_q_rb);
|
|
let x_offset = x_offset.add(&g_lb.unsqueeze(1)?.broadcast_mul(&x_q_lb)?)?;
|
|
// drop(g_lb);
|
|
// drop(x_q_lb);
|
|
let x_offset = x_offset.add(&g_rt.unsqueeze(1)?.broadcast_mul(&x_q_rt)?)?;
|
|
// drop(g_rt);
|
|
// drop(x_q_rt);
|
|
// (bs, n, h, w) -> (bs, h, w, n) -> (bs, 1, h, w, n)
|
|
let m = modulator.permute((0, 2, 3, 1))?.unsqueeze(1)?;
|
|
let x_offset = x_offset.to_device(m.device())?.broadcast_mul(&m)?;
|
|
let x_offset = self.reshape_x_offset(&x_offset, self.ks)?;
|
|
let xs = self.regular_conv.forward(&x_offset)?;
|
|
Ok(xs)
|
|
}
|
|
fn reshape_x_offset(&self, xs: &Tensor, ks: usize) -> Result<Tensor> {
|
|
let (b, c, h, w, _) = xs.dims5()?;
|
|
|
|
let xs = xs.reshape((b, c, h, w, ks, ks))?;
|
|
let xs = xs.permute((0, 1, 2, 4, 3, 5))?;
|
|
let xs = xs.reshape((b, c, h * ks, w * ks))?;
|
|
let x_offset = xs.contiguous()?;
|
|
|
|
Ok(x_offset)
|
|
}
|
|
|
|
fn get_x_q(&self, xs: &Tensor, q_x: &Tensor, q_y: &Tensor) -> Result<Tensor> {
|
|
let (b, h, w, n) = q_x.dims4()?;
|
|
let padded_w = xs.dim(3)?;
|
|
let c = xs.dim(1)?;
|
|
|
|
// 展平输入: (b, c, H, W) -> (b, c, H*W)
|
|
let xs_flat = xs.flatten_from(2)?;
|
|
|
|
// 计算索引
|
|
let index = q_x.affine(padded_w as f64, 0.0)?.add(q_y)?; // (b, h, w, n)
|
|
|
|
// 扩展维度以匹配通道数
|
|
let index = index
|
|
.unsqueeze(1)?
|
|
.expand((b, c, h, w, n))?
|
|
.flatten_from(2)?; // (b, c, h*w*n)
|
|
|
|
// 收集特征
|
|
let xs = xs_flat.gather(&index, 2)?.reshape((b, c, h, w, n))?;
|
|
Ok(xs)
|
|
}
|
|
|
|
fn get_p_n(&self, n: usize, dtype: DType, device: &Device) -> Result<Tensor> {
|
|
let ks = self.ks as f32;
|
|
let range = Tensor::arange_step(-(ks - 1.0) / 2.0, (ks - 1.0) / 2.0 + 1.0, 1.0, device)?;
|
|
// 假设 ks=3
|
|
// [(-1, -1), (-1, 0), (-1, 1)
|
|
// (0, -1), (0, 0), (0, 1)
|
|
// (1, -1), (1, 0), (1, 1)]
|
|
// range: [-1, 0, 1]
|
|
// unsqueeze(1) -> [[-1], [0], [1]]
|
|
// broadcase_as(3, 3) -> [[-1, -1, -1], [0, 0, 0], [1, 1, 1]]
|
|
// flatten_all -> [-1, -1, -1, 0, 0, 0, 1, 1, 1]
|
|
let p_n_x = range
|
|
.unsqueeze(1)?
|
|
.broadcast_as((self.ks, self.ks))?
|
|
.flatten_all()?;
|
|
// range: [-1, 0, 1]
|
|
// unsqueeze(0) -> [[-1, 0, 1]]
|
|
// broadcase_as(3, 3) -> [[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]]
|
|
// flatten_all -> [-1, 0, 1, -1, 0, 1, -1, 0, 1]
|
|
let p_n_y = range
|
|
.unsqueeze(0)?
|
|
.broadcast_as((self.ks, self.ks))?
|
|
.flatten_all()?;
|
|
let p = Tensor::cat(&[p_n_x, p_n_y], 0)?
|
|
.reshape((1, 2 * n, 1, 1))?
|
|
.to_dtype(dtype)?
|
|
.contiguous()?;
|
|
Ok(p)
|
|
}
|
|
|
|
fn get_p_0(
|
|
&self,
|
|
h: usize,
|
|
w: usize,
|
|
n: usize,
|
|
dtype: DType,
|
|
device: &Device,
|
|
) -> Result<Tensor> {
|
|
// 假设 in featuremap h=w=5, padding=1, hp=wp=7
|
|
// out featuremap h=w=5,
|
|
// padding 后的 in featuremap
|
|
// [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6)
|
|
// (1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (1, 5), (1, 6)
|
|
// (2, 0), (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6)
|
|
// (3, 0), (3, 1), (3, 2), (3, 3), (3, 4), (3, 5), (3, 6)
|
|
// (4, 0), (4, 1), (4, 2), (4, 3), (4, 4), (4, 5), (4, 6)
|
|
// (5, 0), (5, 1), (5, 2), (5, 3), (5, 4), (5, 5), (5, 6)
|
|
// (6, 0), (6, 1), (6, 2), (6, 3), (6, 4), (6, 5), (6, 6)
|
|
let start = self.padding as f32;
|
|
// let start = 0.0f32;
|
|
let p_0_x = Tensor::arange_step(
|
|
start,
|
|
start + h as f32 * self.stride as f32,
|
|
self.stride as f32,
|
|
device,
|
|
)?
|
|
.unsqueeze(1)?
|
|
.broadcast_as((h, w))?
|
|
.reshape((1, 1, h, w))?
|
|
.repeat((1, n, 1, 1))?;
|
|
let p_0_y = Tensor::arange_step(
|
|
start,
|
|
start + w as f32 * self.stride as f32,
|
|
self.stride as f32,
|
|
device,
|
|
)?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((h, w))?
|
|
.reshape((1, 1, h, w))?
|
|
.repeat((1, n, 1, 1))?;
|
|
let p_0 = Tensor::cat(&[p_0_x, p_0_y], 1)?
|
|
.to_dtype(dtype)?
|
|
.contiguous()?;
|
|
Ok(p_0)
|
|
}
|
|
fn get_p(&self, offset: &Tensor) -> Result<Tensor> {
|
|
let (_, n, h, w) = offset.dims4()?;
|
|
let n = n / 2;
|
|
// (1, 2n, 1, 1)
|
|
let p_n = self.get_p_n(n, offset.dtype(), offset.device())?;
|
|
// (1, 2n, h, w)
|
|
let p_0 = self.get_p_0(h, w, n, offset.dtype(), offset.device())?;
|
|
let p = p_0
|
|
.broadcast_add(&p_n)?
|
|
.broadcast_add(offset)?
|
|
.contiguous()?;
|
|
Ok(p)
|
|
}
|
|
}
|
|
|
|
struct _ASPPModuleDeformable {
|
|
atrous_conv: DeformableConv2d,
|
|
bn: BatchNorm,
|
|
}
|
|
|
|
impl _ASPPModuleDeformable {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
in_c: usize,
|
|
out_c: usize,
|
|
kernel_size: usize,
|
|
padding: usize,
|
|
) -> Result<Self> {
|
|
let atrous_conv = DeformableConv2d::new(
|
|
vb.pp("atrous_conv"),
|
|
in_c,
|
|
out_c,
|
|
kernel_size,
|
|
1,
|
|
padding,
|
|
false,
|
|
)?;
|
|
let bn = get_batch_norm(vb.pp("bn"), 1e-5, out_c, true)?;
|
|
Ok(Self { atrous_conv, bn })
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let xs = self.atrous_conv.forward(xs)?;
|
|
let xs = self.bn.forward_t(&xs, false)?.relu()?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
struct ASPPDeformable {
|
|
aspp1: _ASPPModuleDeformable,
|
|
// aspp_deforms: Vec<_ASPPModuleDeformable>,
|
|
aspp_deforms_0: _ASPPModuleDeformable,
|
|
aspp_deforms_1: _ASPPModuleDeformable,
|
|
aspp_deforms_2: _ASPPModuleDeformable,
|
|
// avgpool2d + conv2d + BatchNorm2d + relu
|
|
global_avg_pool_1: Conv2d,
|
|
global_avg_pool_2: BatchNorm,
|
|
conv1: Conv2d,
|
|
bn1: BatchNorm,
|
|
}
|
|
|
|
impl ASPPDeformable {
|
|
pub fn new(
|
|
vb: VarBuilder,
|
|
in_c: usize,
|
|
out_c: usize,
|
|
parallel_block_sizes: Vec<usize>,
|
|
) -> Result<Self> {
|
|
let in_channelster = 256;
|
|
let aspp1 = _ASPPModuleDeformable::new(vb.pp("aspp1"), in_c, in_channelster, 1, 0)?;
|
|
let vb_aspp_deforms = vb.pp("aspp_deforms");
|
|
let aspp_deforms_0 = _ASPPModuleDeformable::new(
|
|
vb_aspp_deforms.pp(0),
|
|
in_c,
|
|
in_channelster,
|
|
parallel_block_sizes[0],
|
|
parallel_block_sizes[0] / 2,
|
|
)?;
|
|
let aspp_deforms_1 = _ASPPModuleDeformable::new(
|
|
vb_aspp_deforms.pp(1),
|
|
in_c,
|
|
in_channelster,
|
|
parallel_block_sizes[1],
|
|
parallel_block_sizes[1] / 2,
|
|
)?;
|
|
let aspp_deforms_2 = _ASPPModuleDeformable::new(
|
|
vb_aspp_deforms.pp(2),
|
|
in_c,
|
|
in_channelster,
|
|
parallel_block_sizes[2],
|
|
parallel_block_sizes[2] / 2,
|
|
)?;
|
|
let global_avg_pool_1 = get_conv2d(
|
|
vb.pp("global_avg_pool.1"),
|
|
in_c,
|
|
in_channelster,
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
false,
|
|
)?;
|
|
let global_avg_pool_2 =
|
|
get_batch_norm(vb.pp("global_avg_pool.2"), 1e-5, in_channelster, true)?;
|
|
let conv1 = get_conv2d(
|
|
vb.pp("conv1"),
|
|
in_channelster * (2 + parallel_block_sizes.len()),
|
|
out_c,
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
false,
|
|
)?;
|
|
let bn1 = get_batch_norm(vb.pp("bn1"), 1e-5, out_c, true)?;
|
|
Ok(Self {
|
|
aspp1,
|
|
aspp_deforms_0,
|
|
aspp_deforms_1,
|
|
aspp_deforms_2,
|
|
global_avg_pool_1,
|
|
global_avg_pool_2,
|
|
conv1,
|
|
bn1,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let x1 = self.aspp1.forward(xs)?;
|
|
let x_aspp_deforms_0 = self.aspp_deforms_0.forward(xs)?;
|
|
let x_aspp_deforms_1 = self.aspp_deforms_1.forward(xs)?;
|
|
let x_aspp_deforms_2 = self.aspp_deforms_2.forward(xs)?;
|
|
|
|
let (_, _, h, w) = xs.dims4()?;
|
|
assert_eq!(h, w, "avg_pool2d h, w mus be equal");
|
|
let x5 = xs.avg_pool2d(h)?;
|
|
let x5 = self.global_avg_pool_1.forward(&x5)?;
|
|
let x5 = self.global_avg_pool_2.forward_t(&x5, false)?.relu()?;
|
|
let (_, _, h, w) = x1.dims4()?;
|
|
let x5 = interpolate_bilinear(&x5, (h, w), Some(true))?;
|
|
let xs = Tensor::cat(
|
|
&[x1, x_aspp_deforms_0, x_aspp_deforms_1, x_aspp_deforms_2, x5],
|
|
1,
|
|
)?;
|
|
let xs = self.conv1.forward(&xs)?;
|
|
let xs = self.bn1.forward_t(&xs, false)?.relu()?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
struct BasicDecBlk {
|
|
conv_in: Conv2d,
|
|
dec_att: ASPPDeformable,
|
|
conv_out: Conv2d,
|
|
bn_in: BatchNorm,
|
|
bn_out: BatchNorm,
|
|
}
|
|
|
|
impl BasicDecBlk {
|
|
pub fn new(vb: VarBuilder, in_c: usize, out_c: usize) -> Result<Self> {
|
|
let inter_channels = 64;
|
|
let conv_in = get_conv2d(vb.pp("conv_in"), in_c, inter_channels, 3, 1, 1, 1, 1, true)?;
|
|
let dec_att = ASPPDeformable::new(vb.pp("dec_att"), inter_channels, inter_channels, vec![
|
|
1, 3, 7,
|
|
])?;
|
|
let conv_out = get_conv2d(
|
|
vb.pp("conv_out"),
|
|
inter_channels,
|
|
out_c,
|
|
3,
|
|
1,
|
|
1,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let bn_in = get_batch_norm(vb.pp("bn_in"), 1e-5, inter_channels, true)?;
|
|
let bn_out = get_batch_norm(vb.pp("bn_out"), 1e-5, out_c, true)?;
|
|
Ok(Self {
|
|
conv_in,
|
|
dec_att,
|
|
conv_out,
|
|
bn_in,
|
|
bn_out,
|
|
})
|
|
}
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let xs = self.conv_in.forward(xs)?;
|
|
let xs = self.bn_in.forward_t(&xs, false)?;
|
|
let xs = xs.relu()?;
|
|
let xs = self.dec_att.forward(&xs)?;
|
|
let xs = self.conv_out.forward(&xs)?;
|
|
let xs = self.bn_out.forward_t(&xs, false)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
struct SimpleConvs {
|
|
conv1: Conv2d,
|
|
conv_out: Conv2d,
|
|
}
|
|
|
|
impl SimpleConvs {
|
|
pub fn new(vb: VarBuilder, in_c: usize, out_c: usize, inter_c: usize) -> Result<Self> {
|
|
// inter_c = 64
|
|
let conv1 = get_conv2d(vb.pp("conv1"), in_c, inter_c, 3, 1, 1, 1, 1, true)?;
|
|
let conv_out = get_conv2d(vb.pp("conv_out"), inter_c, out_c, 3, 1, 1, 1, 1, true)?;
|
|
Ok(Self { conv1, conv_out })
|
|
}
|
|
|
|
pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
|
let x = self.conv1.forward(x)?;
|
|
let x = self.conv_out.forward(&x)?;
|
|
Ok(x)
|
|
}
|
|
}
|
|
|
|
struct Decoder {
|
|
ipt_blk5: SimpleConvs,
|
|
ipt_blk4: SimpleConvs,
|
|
ipt_blk3: SimpleConvs,
|
|
ipt_blk2: SimpleConvs,
|
|
ipt_blk1: SimpleConvs,
|
|
decoder_block4: BasicDecBlk,
|
|
decoder_block3: BasicDecBlk,
|
|
decoder_block2: BasicDecBlk,
|
|
decoder_block1: BasicDecBlk,
|
|
conv_out1: Conv2d,
|
|
// BasicLatBlk : conv, kernel_size=1
|
|
lateral_block4: Conv2d,
|
|
lateral_block3: Conv2d,
|
|
lateral_block2: Conv2d,
|
|
// conv_ms_spvn_4: Conv2d,
|
|
// conv_ms_spvn_3: Conv2d,
|
|
// conv_ms_spvn_2: Conv2d,
|
|
gdt_convs_4: Conv2dWithBN,
|
|
gdt_convs_3: Conv2dWithBN,
|
|
gdt_convs_2: Conv2dWithBN,
|
|
gdt_convs_attn_4: Conv2d,
|
|
gdt_convs_attn_3: Conv2d,
|
|
gdt_convs_attn_2: Conv2d,
|
|
}
|
|
|
|
impl Decoder {
|
|
pub fn new(vb: VarBuilder, channels: Vec<usize>) -> Result<Self> {
|
|
let ic = 64;
|
|
let ipt_blk5 =
|
|
SimpleConvs::new(vb.pp("ipt_blk5"), 2usize.pow(10) * 3, channels[0] / 8, ic)?;
|
|
let ipt_blk4 = SimpleConvs::new(vb.pp("ipt_blk4"), 2usize.pow(8) * 3, channels[0] / 8, ic)?;
|
|
let ipt_blk3 = SimpleConvs::new(vb.pp("ipt_blk3"), 2usize.pow(6) * 3, channels[1] / 8, ic)?;
|
|
let ipt_blk2 = SimpleConvs::new(vb.pp("ipt_blk2"), 2usize.pow(4) * 3, channels[2] / 8, ic)?;
|
|
let ipt_blk1 = SimpleConvs::new(vb.pp("ipt_blk1"), 3, channels[3] / 8, ic)?;
|
|
|
|
let decoder_block4 = BasicDecBlk::new(
|
|
vb.pp("decoder_block4"),
|
|
channels[0] + channels[0] / 8,
|
|
channels[1],
|
|
)?;
|
|
let decoder_block3 = BasicDecBlk::new(
|
|
vb.pp("decoder_block3"),
|
|
channels[1] + channels[0] / 8,
|
|
channels[2],
|
|
)?;
|
|
let decoder_block2 = BasicDecBlk::new(
|
|
vb.pp("decoder_block2"),
|
|
channels[2] + channels[1] / 8,
|
|
channels[3],
|
|
)?;
|
|
let decoder_block1 = BasicDecBlk::new(
|
|
vb.pp("decoder_block1"),
|
|
channels[3] + channels[2] / 8,
|
|
channels[3] / 2,
|
|
)?;
|
|
|
|
let conv_out1 = get_conv2d(
|
|
vb.pp("conv_out1.0"),
|
|
channels[3] / 2 + channels[3] / 8,
|
|
1,
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let lateral_block4 = get_conv2d(
|
|
vb.pp("lateral_block4.conv"),
|
|
channels[1],
|
|
channels[1],
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let lateral_block3 = get_conv2d(
|
|
vb.pp("lateral_block3.conv"),
|
|
channels[2],
|
|
channels[2],
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let lateral_block2 = get_conv2d(
|
|
vb.pp("lateral_block2.conv"),
|
|
channels[3],
|
|
channels[3],
|
|
1,
|
|
0,
|
|
1,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
|
|
// let conv_ms_spvn_4 =
|
|
// get_conv2d(vb.pp("conv_ms_spvn_4"), channels[1], 1, 1, 0, 1, 1, 1, true)?;
|
|
// let conv_ms_spvn_3 =
|
|
// get_conv2d(vb.pp("conv_ms_spvn_3"), channels[2], 1, 1, 0, 1, 1, 1, true)?;
|
|
// let conv_ms_spvn_2 =
|
|
// get_conv2d(vb.pp("conv_ms_spvn_2"), channels[3], 1, 1, 0, 1, 1, 1, true)?;
|
|
let n = 16usize;
|
|
let gdt_convs_4 =
|
|
Conv2dWithBN::new(vb.pp("gdt_convs_4"), channels[1], n, 3, 1, 1, true, true)?;
|
|
let gdt_convs_3 =
|
|
Conv2dWithBN::new(vb.pp("gdt_convs_3"), channels[2], n, 3, 1, 1, true, true)?;
|
|
let gdt_convs_2 =
|
|
Conv2dWithBN::new(vb.pp("gdt_convs_2"), channels[3], n, 3, 1, 1, true, true)?;
|
|
|
|
let gdt_convs_attn_4 = get_conv2d(vb.pp("gdt_convs_attn_4.0"), n, 1, 1, 0, 1, 1, 1, true)?;
|
|
let gdt_convs_attn_3 = get_conv2d(vb.pp("gdt_convs_attn_3.0"), n, 1, 1, 0, 1, 1, 1, true)?;
|
|
let gdt_convs_attn_2 = get_conv2d(vb.pp("gdt_convs_attn_2.0"), n, 1, 1, 0, 1, 1, 1, true)?;
|
|
Ok(Self {
|
|
ipt_blk5,
|
|
ipt_blk4,
|
|
ipt_blk3,
|
|
ipt_blk2,
|
|
ipt_blk1,
|
|
decoder_block4,
|
|
decoder_block3,
|
|
decoder_block2,
|
|
decoder_block1,
|
|
conv_out1,
|
|
lateral_block4,
|
|
lateral_block3,
|
|
lateral_block2,
|
|
// conv_ms_spvn_4,
|
|
// conv_ms_spvn_3,
|
|
// conv_ms_spvn_2,
|
|
gdt_convs_4,
|
|
gdt_convs_3,
|
|
gdt_convs_2,
|
|
gdt_convs_attn_4,
|
|
gdt_convs_attn_3,
|
|
gdt_convs_attn_2,
|
|
})
|
|
}
|
|
|
|
pub fn get_patches_batch(&self, x: &Tensor, p: &Tensor) -> Result<Tensor> {
|
|
let (_, _, h, w) = p.dims4()?;
|
|
let mut patches_batch = vec![];
|
|
for idx in 0..x.dim(0)? {
|
|
let x_i = x.i(idx)?.unsqueeze(0)?; // 保证bs的维度存在
|
|
let columns_x = split_tensor_with_size(&x_i, w, D::Minus1)?;
|
|
let mut patches_x = vec![];
|
|
for col_x in columns_x {
|
|
let pat_x = split_tensor_with_size(&col_x, h, D::Minus2)?;
|
|
patches_x.extend_from_slice(&pat_x);
|
|
}
|
|
let patch_sample = Tensor::cat(&patches_x, 1)?;
|
|
patches_batch.push(patch_sample);
|
|
}
|
|
let patch = Tensor::cat(&patches_batch, 0)?;
|
|
Ok(patch)
|
|
}
|
|
|
|
pub fn forward(&self, features: Vec<&Tensor>) -> Result<Tensor> {
|
|
let [x, x1, x2, x3, x4] = features[..] else {
|
|
return Err(anyhow!(format!(
|
|
"swintransformer output exactly 3 elements"
|
|
)));
|
|
};
|
|
// let mut outs = vec![];
|
|
let patches_batch = self.get_patches_batch(x, x4)?;
|
|
let (_, _, x4_h, x4_w) = x4.dims4()?;
|
|
let patches_batch = interpolate_bilinear(&patches_batch, (x4_h, x4_w), Some(true))?;
|
|
let ipt_blk5_out = self.ipt_blk5.forward(&patches_batch)?;
|
|
let x4 = Tensor::cat(&[x4, &ipt_blk5_out], 1)?;
|
|
let p4 = self.decoder_block4.forward(&x4)?;
|
|
// let m4 = self.conv_ms_spvn_4.forward(&p4)?;
|
|
// outs.push(m4);
|
|
let p4_gdt = self.gdt_convs_4.forward(&p4)?;
|
|
let gdt_attn_4 = sigmoid(&self.gdt_convs_attn_4.forward(&p4_gdt)?)?;
|
|
let p4 = p4.broadcast_mul(&gdt_attn_4)?;
|
|
|
|
let (_, _, x3_h, x3_w) = x3.dims4()?;
|
|
let p4_inter = interpolate_bilinear(&p4, (x3_h, x3_w), Some(true))?;
|
|
let p3_ = self.lateral_block4.forward(x3)?;
|
|
let p3_ = p4_inter.add(&p3_)?;
|
|
let patches_batch = self.get_patches_batch(x, &p3_)?;
|
|
let patches_batch = interpolate_bilinear(&patches_batch, (x3_h, x3_w), Some(true))?;
|
|
let ipt_blk4_out = self.ipt_blk4.forward(&patches_batch)?;
|
|
let p3_ = Tensor::cat(&[p3_, ipt_blk4_out], 1)?;
|
|
let p3 = self.decoder_block3.forward(&p3_)?;
|
|
// let m3 = self.conv_ms_spvn_3.forward(&p3)?;
|
|
// outs.push(m3);
|
|
let p3_gdt = self.gdt_convs_3.forward(&p3)?;
|
|
let gdt_attn_3 = sigmoid(&self.gdt_convs_attn_3.forward(&p3_gdt)?)?;
|
|
let p3 = p3.broadcast_mul(&gdt_attn_3)?;
|
|
|
|
let (_, _, x2_h, x2_w) = x2.dims4()?;
|
|
let p3_inter = interpolate_bilinear(&p3, (x2_h, x2_w), Some(true))?;
|
|
let p2_ = self.lateral_block3.forward(x2)?;
|
|
let p2_ = p3_inter.add(&p2_)?;
|
|
let patches_batch = self.get_patches_batch(x, &p2_)?;
|
|
let patches_batch = interpolate_bilinear(&patches_batch, (x2_h, x2_w), Some(true))?;
|
|
let ipt_blk3_out = self.ipt_blk3.forward(&patches_batch)?;
|
|
let p2_ = Tensor::cat(&[p2_, ipt_blk3_out], 1)?;
|
|
let p2 = self.decoder_block2.forward(&p2_)?;
|
|
// let m2 = self.conv_ms_spvn_2.forward(&p2)?;
|
|
// outs.push(m2);
|
|
let p2_gdt = self.gdt_convs_2.forward(&p2)?;
|
|
let gdt_attn_2 = sigmoid(&self.gdt_convs_attn_2.forward(&p2_gdt)?)?;
|
|
let p2 = p2.broadcast_mul(&gdt_attn_2)?;
|
|
|
|
let (_, _, x1_h, x1_w) = x1.dims4()?;
|
|
let p2_inter = interpolate_bilinear(&p2, (x1_h, x1_w), Some(true))?;
|
|
let p1_ = self.lateral_block2.forward(x1)?;
|
|
let p1_ = p2_inter.add(&p1_)?;
|
|
let patches_batch = self.get_patches_batch(x, &p1_)?;
|
|
let patches_batch = interpolate_bilinear(&patches_batch, (x1_h, x1_w), Some(true))?;
|
|
let ipt_blk2_out = self.ipt_blk2.forward(&patches_batch)?;
|
|
let p1_ = Tensor::cat(&[p1_, ipt_blk2_out], 1)?;
|
|
let p1_ = self.decoder_block1.forward(&p1_)?;
|
|
|
|
let (_, _, x_h, x_w) = x.dims4()?;
|
|
let p1_ = interpolate_bilinear(&p1_, (x_h, x_w), Some(true))?;
|
|
// let patches_batch = self.get_patches_batch(x, &p1_)?;
|
|
// let patches_batch = interpolate_bilinear(&patches_batch, (x_h, x_w), Some(true))?;
|
|
let ipt_blk1_out = self.ipt_blk1.forward(x)?;
|
|
let p1_ = Tensor::cat(&[p1_, ipt_blk1_out], 1)?;
|
|
let p1_out = self.conv_out1.forward(&p1_)?;
|
|
let out = sigmoid(&p1_out)?;
|
|
// outs.push(p1_out);
|
|
Ok(out)
|
|
}
|
|
}
|
|
|
|
pub struct BiRefNet {
|
|
bb: SwinTransformer,
|
|
squeeze_module_0: BasicDecBlk,
|
|
decoder: Decoder,
|
|
}
|
|
impl BiRefNet {
|
|
pub fn new(vb: VarBuilder) -> Result<Self> {
|
|
let bb = SwinTransformer::new(
|
|
vb.pp("bb"),
|
|
4,
|
|
3,
|
|
192,
|
|
vec![2, 2, 18, 2],
|
|
vec![6, 12, 24, 48],
|
|
12,
|
|
4.0,
|
|
true,
|
|
true,
|
|
vec![0, 1, 2, 3],
|
|
)?;
|
|
let channels = vec![3072, 1536, 768, 384];
|
|
let in_c = channels.iter().sum();
|
|
let squeeze_module_0 = BasicDecBlk::new(vb.pp("squeeze_module.0"), in_c, channels[0])?;
|
|
let decoder = Decoder::new(vb.pp("decoder"), channels)?;
|
|
Ok(Self {
|
|
bb,
|
|
squeeze_module_0,
|
|
decoder,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
|
let [ref x1, ref x2, ref x3, ref x4] = self.bb.forward(xs)?[..] else {
|
|
return Err(anyhow!(format!(
|
|
"swintransformer output exactly 3 elements"
|
|
)));
|
|
};
|
|
|
|
let (_, _, h, w) = xs.dims4()?;
|
|
let cat_xs = interpolate_bilinear(xs, (h / 2, w / 2), Some(true))?;
|
|
let [ref x1_, ref x2_, ref x3_, ref x4_] = self.bb.forward(&cat_xs)?[..] else {
|
|
return Err(anyhow!(format!(
|
|
"swintransformer output exactly 3 elements"
|
|
)));
|
|
};
|
|
let (_, _, x1_h, x1_w) = x1.dims4()?;
|
|
let x1_ = interpolate_bilinear(x1_, (x1_h, x1_w), Some(true))?;
|
|
let x1 = Tensor::cat(&[x1, &x1_], 1)?;
|
|
let (_, _, x2_h, x2_w) = x2.dims4()?;
|
|
let x2_ = interpolate_bilinear(x2_, (x2_h, x2_w), Some(true))?;
|
|
let x2 = Tensor::cat(&[x2, &x2_], 1)?;
|
|
let (_, _, x3_h, x3_w) = x3.dims4()?;
|
|
let x3_ = interpolate_bilinear(x3_, (x3_h, x3_w), Some(true))?;
|
|
let x3 = Tensor::cat(&[x3, &x3_], 1)?;
|
|
let (_, _, x4_h, x4_w) = x4.dims4()?;
|
|
let x4_ = interpolate_bilinear(x4_, (x4_h, x4_w), Some(true))?;
|
|
let x4 = Tensor::cat(&[x4, &x4_], 1)?;
|
|
let x1_resize = interpolate_bilinear(&x1, (x4_h, x4_w), Some(true))?;
|
|
let x2_resize = interpolate_bilinear(&x2, (x4_h, x4_w), Some(true))?;
|
|
let x3_resize = interpolate_bilinear(&x3, (x4_h, x4_w), Some(true))?;
|
|
let x4 = Tensor::cat(&[x1_resize, x2_resize, x3_resize, x4], 1)?;
|
|
|
|
let x4 = self.squeeze_module_0.forward(&x4)?;
|
|
|
|
let features = vec![xs, &x1, &x2, &x3, &x4];
|
|
let output = self.decoder.forward(features)?;
|
|
Ok(output)
|
|
}
|
|
}
|