Skip to content
Projects
Groups
Snippets
Help
Loading...
Help
Contribute to GitLab
Sign in / Register
Toggle navigation
S
stable-diffusion-webui
Project
Project
Details
Activity
Cycle Analytics
Repository
Repository
Files
Commits
Branches
Tags
Contributors
Graph
Compare
Charts
Issues
0
Issues
0
List
Board
Labels
Milestones
Merge Requests
0
Merge Requests
0
CI / CD
CI / CD
Pipelines
Jobs
Schedules
Charts
Wiki
Wiki
Snippets
Snippets
Members
Members
Collapse sidebar
Close sidebar
Activity
Graph
Charts
Create a new issue
Jobs
Commits
Issue Boards
Open sidebar
Administrator
stable-diffusion-webui
Commits
bff8d0ce
Commit
bff8d0ce
authored
Sep 26, 2022
by
d8ahazard
Browse files
Options
Browse Files
Download
Email Patches
Plain Diff
Fix swinir arch name
parent
78d6aef3
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
867 additions
and
867 deletions
+867
-867
swinir_model_arch.py
modules/swinir_model_arch.py
+867
-867
No files found.
modules/swinir_arch.py
→
modules/swinir_
model_
arch.py
View file @
bff8d0ce
# -----------------------------------------------------------------------------------
# -----------------------------------------------------------------------------------
# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257
# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257
# Originally Written by Ze Liu, Modified by Jingyun Liang.
# Originally Written by Ze Liu, Modified by Jingyun Liang.
# -----------------------------------------------------------------------------------
# -----------------------------------------------------------------------------------
import
math
import
math
import
torch
import
torch
import
torch.nn
as
nn
import
torch.nn
as
nn
import
torch.nn.functional
as
F
import
torch.nn.functional
as
F
import
torch.utils.checkpoint
as
checkpoint
import
torch.utils.checkpoint
as
checkpoint
from
timm.models.layers
import
DropPath
,
to_2tuple
,
trunc_normal_
from
timm.models.layers
import
DropPath
,
to_2tuple
,
trunc_normal_
class
Mlp
(
nn
.
Module
):
class
Mlp
(
nn
.
Module
):
def
__init__
(
self
,
in_features
,
hidden_features
=
None
,
out_features
=
None
,
act_layer
=
nn
.
GELU
,
drop
=
0.
):
def
__init__
(
self
,
in_features
,
hidden_features
=
None
,
out_features
=
None
,
act_layer
=
nn
.
GELU
,
drop
=
0.
):
super
()
.
__init__
()
super
()
.
__init__
()
out_features
=
out_features
or
in_features
out_features
=
out_features
or
in_features
hidden_features
=
hidden_features
or
in_features
hidden_features
=
hidden_features
or
in_features
self
.
fc1
=
nn
.
Linear
(
in_features
,
hidden_features
)
self
.
fc1
=
nn
.
Linear
(
in_features
,
hidden_features
)
self
.
act
=
act_layer
()
self
.
act
=
act_layer
()
self
.
fc2
=
nn
.
Linear
(
hidden_features
,
out_features
)
self
.
fc2
=
nn
.
Linear
(
hidden_features
,
out_features
)
self
.
drop
=
nn
.
Dropout
(
drop
)
self
.
drop
=
nn
.
Dropout
(
drop
)
def
forward
(
self
,
x
):
def
forward
(
self
,
x
):
x
=
self
.
fc1
(
x
)
x
=
self
.
fc1
(
x
)
x
=
self
.
act
(
x
)
x
=
self
.
act
(
x
)
x
=
self
.
drop
(
x
)
x
=
self
.
drop
(
x
)
x
=
self
.
fc2
(
x
)
x
=
self
.
fc2
(
x
)
x
=
self
.
drop
(
x
)
x
=
self
.
drop
(
x
)
return
x
return
x
def
window_partition
(
x
,
window_size
):
def
window_partition
(
x
,
window_size
):
"""
"""
Args:
Args:
x: (B, H, W, C)
x: (B, H, W, C)
window_size (int): window size
window_size (int): window size
Returns:
Returns:
windows: (num_windows*B, window_size, window_size, C)
windows: (num_windows*B, window_size, window_size, C)
"""
"""
B
,
H
,
W
,
C
=
x
.
shape
B
,
H
,
W
,
C
=
x
.
shape
x
=
x
.
view
(
B
,
H
//
window_size
,
window_size
,
W
//
window_size
,
window_size
,
C
)
x
=
x
.
view
(
B
,
H
//
window_size
,
window_size
,
W
//
window_size
,
window_size
,
C
)
windows
=
x
.
permute
(
0
,
1
,
3
,
2
,
4
,
5
)
.
contiguous
()
.
view
(
-
1
,
window_size
,
window_size
,
C
)
windows
=
x
.
permute
(
0
,
1
,
3
,
2
,
4
,
5
)
.
contiguous
()
.
view
(
-
1
,
window_size
,
window_size
,
C
)
return
windows
return
windows
def
window_reverse
(
windows
,
window_size
,
H
,
W
):
def
window_reverse
(
windows
,
window_size
,
H
,
W
):
"""
"""
Args:
Args:
windows: (num_windows*B, window_size, window_size, C)
windows: (num_windows*B, window_size, window_size, C)
window_size (int): Window size
window_size (int): Window size
H (int): Height of image
H (int): Height of image
W (int): Width of image
W (int): Width of image
Returns:
Returns:
x: (B, H, W, C)
x: (B, H, W, C)
"""
"""
B
=
int
(
windows
.
shape
[
0
]
/
(
H
*
W
/
window_size
/
window_size
))
B
=
int
(
windows
.
shape
[
0
]
/
(
H
*
W
/
window_size
/
window_size
))
x
=
windows
.
view
(
B
,
H
//
window_size
,
W
//
window_size
,
window_size
,
window_size
,
-
1
)
x
=
windows
.
view
(
B
,
H
//
window_size
,
W
//
window_size
,
window_size
,
window_size
,
-
1
)
x
=
x
.
permute
(
0
,
1
,
3
,
2
,
4
,
5
)
.
contiguous
()
.
view
(
B
,
H
,
W
,
-
1
)
x
=
x
.
permute
(
0
,
1
,
3
,
2
,
4
,
5
)
.
contiguous
()
.
view
(
B
,
H
,
W
,
-
1
)
return
x
return
x
class
WindowAttention
(
nn
.
Module
):
class
WindowAttention
(
nn
.
Module
):
r""" Window based multi-head self attention (W-MSA) module with relative position bias.
r""" Window based multi-head self attention (W-MSA) module with relative position bias.
It supports both of shifted and non-shifted window.
It supports both of shifted and non-shifted window.
Args:
Args:
dim (int): Number of input channels.
dim (int): Number of input channels.
window_size (tuple[int]): The height and width of the window.
window_size (tuple[int]): The height and width of the window.
num_heads (int): Number of attention heads.
num_heads (int): Number of attention heads.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set
attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
proj_drop (float, optional): Dropout ratio of output. Default: 0.0
"""
"""
def
__init__
(
self
,
dim
,
window_size
,
num_heads
,
qkv_bias
=
True
,
qk_scale
=
None
,
attn_drop
=
0.
,
proj_drop
=
0.
):
def
__init__
(
self
,
dim
,
window_size
,
num_heads
,
qkv_bias
=
True
,
qk_scale
=
None
,
attn_drop
=
0.
,
proj_drop
=
0.
):
super
()
.
__init__
()
super
()
.
__init__
()
self
.
dim
=
dim
self
.
dim
=
dim
self
.
window_size
=
window_size
# Wh, Ww
self
.
window_size
=
window_size
# Wh, Ww
self
.
num_heads
=
num_heads
self
.
num_heads
=
num_heads
head_dim
=
dim
//
num_heads
head_dim
=
dim
//
num_heads
self
.
scale
=
qk_scale
or
head_dim
**
-
0.5
self
.
scale
=
qk_scale
or
head_dim
**
-
0.5
# define a parameter table of relative position bias
# define a parameter table of relative position bias
self
.
relative_position_bias_table
=
nn
.
Parameter
(
self
.
relative_position_bias_table
=
nn
.
Parameter
(
torch
.
zeros
((
2
*
window_size
[
0
]
-
1
)
*
(
2
*
window_size
[
1
]
-
1
),
num_heads
))
# 2*Wh-1 * 2*Ww-1, nH
torch
.
zeros
((
2
*
window_size
[
0
]
-
1
)
*
(
2
*
window_size
[
1
]
-
1
),
num_heads
))
# 2*Wh-1 * 2*Ww-1, nH
# get pair-wise relative position index for each token inside the window
# get pair-wise relative position index for each token inside the window
coords_h
=
torch
.
arange
(
self
.
window_size
[
0
])
coords_h
=
torch
.
arange
(
self
.
window_size
[
0
])
coords_w
=
torch
.
arange
(
self
.
window_size
[
1
])
coords_w
=
torch
.
arange
(
self
.
window_size
[
1
])
coords
=
torch
.
stack
(
torch
.
meshgrid
([
coords_h
,
coords_w
]))
# 2, Wh, Ww
coords
=
torch
.
stack
(
torch
.
meshgrid
([
coords_h
,
coords_w
]))
# 2, Wh, Ww
coords_flatten
=
torch
.
flatten
(
coords
,
1
)
# 2, Wh*Ww
coords_flatten
=
torch
.
flatten
(
coords
,
1
)
# 2, Wh*Ww
relative_coords
=
coords_flatten
[:,
:,
None
]
-
coords_flatten
[:,
None
,
:]
# 2, Wh*Ww, Wh*Ww
relative_coords
=
coords_flatten
[:,
:,
None
]
-
coords_flatten
[:,
None
,
:]
# 2, Wh*Ww, Wh*Ww
relative_coords
=
relative_coords
.
permute
(
1
,
2
,
0
)
.
contiguous
()
# Wh*Ww, Wh*Ww, 2
relative_coords
=
relative_coords
.
permute
(
1
,
2
,
0
)
.
contiguous
()
# Wh*Ww, Wh*Ww, 2
relative_coords
[:,
:,
0
]
+=
self
.
window_size
[
0
]
-
1
# shift to start from 0
relative_coords
[:,
:,
0
]
+=
self
.
window_size
[
0
]
-
1
# shift to start from 0
relative_coords
[:,
:,
1
]
+=
self
.
window_size
[
1
]
-
1
relative_coords
[:,
:,
1
]
+=
self
.
window_size
[
1
]
-
1
relative_coords
[:,
:,
0
]
*=
2
*
self
.
window_size
[
1
]
-
1
relative_coords
[:,
:,
0
]
*=
2
*
self
.
window_size
[
1
]
-
1
relative_position_index
=
relative_coords
.
sum
(
-
1
)
# Wh*Ww, Wh*Ww
relative_position_index
=
relative_coords
.
sum
(
-
1
)
# Wh*Ww, Wh*Ww
self
.
register_buffer
(
"relative_position_index"
,
relative_position_index
)
self
.
register_buffer
(
"relative_position_index"
,
relative_position_index
)
self
.
qkv
=
nn
.
Linear
(
dim
,
dim
*
3
,
bias
=
qkv_bias
)
self
.
qkv
=
nn
.
Linear
(
dim
,
dim
*
3
,
bias
=
qkv_bias
)
self
.
attn_drop
=
nn
.
Dropout
(
attn_drop
)
self
.
attn_drop
=
nn
.
Dropout
(
attn_drop
)
self
.
proj
=
nn
.
Linear
(
dim
,
dim
)
self
.
proj
=
nn
.
Linear
(
dim
,
dim
)
self
.
proj_drop
=
nn
.
Dropout
(
proj_drop
)
self
.
proj_drop
=
nn
.
Dropout
(
proj_drop
)
trunc_normal_
(
self
.
relative_position_bias_table
,
std
=
.02
)
trunc_normal_
(
self
.
relative_position_bias_table
,
std
=
.02
)
self
.
softmax
=
nn
.
Softmax
(
dim
=-
1
)
self
.
softmax
=
nn
.
Softmax
(
dim
=-
1
)
def
forward
(
self
,
x
,
mask
=
None
):
def
forward
(
self
,
x
,
mask
=
None
):
"""
"""
Args:
Args:
x: input features with shape of (num_windows*B, N, C)
x: input features with shape of (num_windows*B, N, C)
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None
"""
"""
B_
,
N
,
C
=
x
.
shape
B_
,
N
,
C
=
x
.
shape
qkv
=
self
.
qkv
(
x
)
.
reshape
(
B_
,
N
,
3
,
self
.
num_heads
,
C
//
self
.
num_heads
)
.
permute
(
2
,
0
,
3
,
1
,
4
)
qkv
=
self
.
qkv
(
x
)
.
reshape
(
B_
,
N
,
3
,
self
.
num_heads
,
C
//
self
.
num_heads
)
.
permute
(
2
,
0
,
3
,
1
,
4
)
q
,
k
,
v
=
qkv
[
0
],
qkv
[
1
],
qkv
[
2
]
# make torchscript happy (cannot use tensor as tuple)
q
,
k
,
v
=
qkv
[
0
],
qkv
[
1
],
qkv
[
2
]
# make torchscript happy (cannot use tensor as tuple)
q
=
q
*
self
.
scale
q
=
q
*
self
.
scale
attn
=
(
q
@
k
.
transpose
(
-
2
,
-
1
))
attn
=
(
q
@
k
.
transpose
(
-
2
,
-
1
))
relative_position_bias
=
self
.
relative_position_bias_table
[
self
.
relative_position_index
.
view
(
-
1
)]
.
view
(
relative_position_bias
=
self
.
relative_position_bias_table
[
self
.
relative_position_index
.
view
(
-
1
)]
.
view
(
self
.
window_size
[
0
]
*
self
.
window_size
[
1
],
self
.
window_size
[
0
]
*
self
.
window_size
[
1
],
-
1
)
# Wh*Ww,Wh*Ww,nH
self
.
window_size
[
0
]
*
self
.
window_size
[
1
],
self
.
window_size
[
0
]
*
self
.
window_size
[
1
],
-
1
)
# Wh*Ww,Wh*Ww,nH
relative_position_bias
=
relative_position_bias
.
permute
(
2
,
0
,
1
)
.
contiguous
()
# nH, Wh*Ww, Wh*Ww
relative_position_bias
=
relative_position_bias
.
permute
(
2
,
0
,
1
)
.
contiguous
()
# nH, Wh*Ww, Wh*Ww
attn
=
attn
+
relative_position_bias
.
unsqueeze
(
0
)
attn
=
attn
+
relative_position_bias
.
unsqueeze
(
0
)
if
mask
is
not
None
:
if
mask
is
not
None
:
nW
=
mask
.
shape
[
0
]
nW
=
mask
.
shape
[
0
]
attn
=
attn
.
view
(
B_
//
nW
,
nW
,
self
.
num_heads
,
N
,
N
)
+
mask
.
unsqueeze
(
1
)
.
unsqueeze
(
0
)
attn
=
attn
.
view
(
B_
//
nW
,
nW
,
self
.
num_heads
,
N
,
N
)
+
mask
.
unsqueeze
(
1
)
.
unsqueeze
(
0
)
attn
=
attn
.
view
(
-
1
,
self
.
num_heads
,
N
,
N
)
attn
=
attn
.
view
(
-
1
,
self
.
num_heads
,
N
,
N
)
attn
=
self
.
softmax
(
attn
)
attn
=
self
.
softmax
(
attn
)
else
:
else
:
attn
=
self
.
softmax
(
attn
)
attn
=
self
.
softmax
(
attn
)
attn
=
self
.
attn_drop
(
attn
)
attn
=
self
.
attn_drop
(
attn
)
x
=
(
attn
@
v
)
.
transpose
(
1
,
2
)
.
reshape
(
B_
,
N
,
C
)
x
=
(
attn
@
v
)
.
transpose
(
1
,
2
)
.
reshape
(
B_
,
N
,
C
)
x
=
self
.
proj
(
x
)
x
=
self
.
proj
(
x
)
x
=
self
.
proj_drop
(
x
)
x
=
self
.
proj_drop
(
x
)
return
x
return
x
def
extra_repr
(
self
)
->
str
:
def
extra_repr
(
self
)
->
str
:
return
f
'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
return
f
'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'
def
flops
(
self
,
N
):
def
flops
(
self
,
N
):
# calculate flops for 1 window with token length of N
# calculate flops for 1 window with token length of N
flops
=
0
flops
=
0
# qkv = self.qkv(x)
# qkv = self.qkv(x)
flops
+=
N
*
self
.
dim
*
3
*
self
.
dim
flops
+=
N
*
self
.
dim
*
3
*
self
.
dim
# attn = (q @ k.transpose(-2, -1))
# attn = (q @ k.transpose(-2, -1))
flops
+=
self
.
num_heads
*
N
*
(
self
.
dim
//
self
.
num_heads
)
*
N
flops
+=
self
.
num_heads
*
N
*
(
self
.
dim
//
self
.
num_heads
)
*
N
# x = (attn @ v)
# x = (attn @ v)
flops
+=
self
.
num_heads
*
N
*
N
*
(
self
.
dim
//
self
.
num_heads
)
flops
+=
self
.
num_heads
*
N
*
N
*
(
self
.
dim
//
self
.
num_heads
)
# x = self.proj(x)
# x = self.proj(x)
flops
+=
N
*
self
.
dim
*
self
.
dim
flops
+=
N
*
self
.
dim
*
self
.
dim
return
flops
return
flops
class
SwinTransformerBlock
(
nn
.
Module
):
class
SwinTransformerBlock
(
nn
.
Module
):
r""" Swin Transformer Block.
r""" Swin Transformer Block.
Args:
Args:
dim (int): Number of input channels.
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resulotion.
input_resolution (tuple[int]): Input resulotion.
num_heads (int): Number of attention heads.
num_heads (int): Number of attention heads.
window_size (int): Window size.
window_size (int): Window size.
shift_size (int): Shift size for SW-MSA.
shift_size (int): Shift size for SW-MSA.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float, optional): Stochastic depth rate. Default: 0.0
drop_path (float, optional): Stochastic depth rate. Default: 0.0
act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
act_layer (nn.Module, optional): Activation layer. Default: nn.GELU
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
"""
"""
def
__init__
(
self
,
dim
,
input_resolution
,
num_heads
,
window_size
=
7
,
shift_size
=
0
,
def
__init__
(
self
,
dim
,
input_resolution
,
num_heads
,
window_size
=
7
,
shift_size
=
0
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop
=
0.
,
attn_drop
=
0.
,
drop_path
=
0.
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop
=
0.
,
attn_drop
=
0.
,
drop_path
=
0.
,
act_layer
=
nn
.
GELU
,
norm_layer
=
nn
.
LayerNorm
):
act_layer
=
nn
.
GELU
,
norm_layer
=
nn
.
LayerNorm
):
super
()
.
__init__
()
super
()
.
__init__
()
self
.
dim
=
dim
self
.
dim
=
dim
self
.
input_resolution
=
input_resolution
self
.
input_resolution
=
input_resolution
self
.
num_heads
=
num_heads
self
.
num_heads
=
num_heads
self
.
window_size
=
window_size
self
.
window_size
=
window_size
self
.
shift_size
=
shift_size
self
.
shift_size
=
shift_size
self
.
mlp_ratio
=
mlp_ratio
self
.
mlp_ratio
=
mlp_ratio
if
min
(
self
.
input_resolution
)
<=
self
.
window_size
:
if
min
(
self
.
input_resolution
)
<=
self
.
window_size
:
# if window size is larger than input resolution, we don't partition windows
# if window size is larger than input resolution, we don't partition windows
self
.
shift_size
=
0
self
.
shift_size
=
0
self
.
window_size
=
min
(
self
.
input_resolution
)
self
.
window_size
=
min
(
self
.
input_resolution
)
assert
0
<=
self
.
shift_size
<
self
.
window_size
,
"shift_size must in 0-window_size"
assert
0
<=
self
.
shift_size
<
self
.
window_size
,
"shift_size must in 0-window_size"
self
.
norm1
=
norm_layer
(
dim
)
self
.
norm1
=
norm_layer
(
dim
)
self
.
attn
=
WindowAttention
(
self
.
attn
=
WindowAttention
(
dim
,
window_size
=
to_2tuple
(
self
.
window_size
),
num_heads
=
num_heads
,
dim
,
window_size
=
to_2tuple
(
self
.
window_size
),
num_heads
=
num_heads
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
attn_drop
=
attn_drop
,
proj_drop
=
drop
)
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
attn_drop
=
attn_drop
,
proj_drop
=
drop
)
self
.
drop_path
=
DropPath
(
drop_path
)
if
drop_path
>
0.
else
nn
.
Identity
()
self
.
drop_path
=
DropPath
(
drop_path
)
if
drop_path
>
0.
else
nn
.
Identity
()
self
.
norm2
=
norm_layer
(
dim
)
self
.
norm2
=
norm_layer
(
dim
)
mlp_hidden_dim
=
int
(
dim
*
mlp_ratio
)
mlp_hidden_dim
=
int
(
dim
*
mlp_ratio
)
self
.
mlp
=
Mlp
(
in_features
=
dim
,
hidden_features
=
mlp_hidden_dim
,
act_layer
=
act_layer
,
drop
=
drop
)
self
.
mlp
=
Mlp
(
in_features
=
dim
,
hidden_features
=
mlp_hidden_dim
,
act_layer
=
act_layer
,
drop
=
drop
)
if
self
.
shift_size
>
0
:
if
self
.
shift_size
>
0
:
attn_mask
=
self
.
calculate_mask
(
self
.
input_resolution
)
attn_mask
=
self
.
calculate_mask
(
self
.
input_resolution
)
else
:
else
:
attn_mask
=
None
attn_mask
=
None
self
.
register_buffer
(
"attn_mask"
,
attn_mask
)
self
.
register_buffer
(
"attn_mask"
,
attn_mask
)
def
calculate_mask
(
self
,
x_size
):
def
calculate_mask
(
self
,
x_size
):
# calculate attention mask for SW-MSA
# calculate attention mask for SW-MSA
H
,
W
=
x_size
H
,
W
=
x_size
img_mask
=
torch
.
zeros
((
1
,
H
,
W
,
1
))
# 1 H W 1
img_mask
=
torch
.
zeros
((
1
,
H
,
W
,
1
))
# 1 H W 1
h_slices
=
(
slice
(
0
,
-
self
.
window_size
),
h_slices
=
(
slice
(
0
,
-
self
.
window_size
),
slice
(
-
self
.
window_size
,
-
self
.
shift_size
),
slice
(
-
self
.
window_size
,
-
self
.
shift_size
),
slice
(
-
self
.
shift_size
,
None
))
slice
(
-
self
.
shift_size
,
None
))
w_slices
=
(
slice
(
0
,
-
self
.
window_size
),
w_slices
=
(
slice
(
0
,
-
self
.
window_size
),
slice
(
-
self
.
window_size
,
-
self
.
shift_size
),
slice
(
-
self
.
window_size
,
-
self
.
shift_size
),
slice
(
-
self
.
shift_size
,
None
))
slice
(
-
self
.
shift_size
,
None
))
cnt
=
0
cnt
=
0
for
h
in
h_slices
:
for
h
in
h_slices
:
for
w
in
w_slices
:
for
w
in
w_slices
:
img_mask
[:,
h
,
w
,
:]
=
cnt
img_mask
[:,
h
,
w
,
:]
=
cnt
cnt
+=
1
cnt
+=
1
mask_windows
=
window_partition
(
img_mask
,
self
.
window_size
)
# nW, window_size, window_size, 1
mask_windows
=
window_partition
(
img_mask
,
self
.
window_size
)
# nW, window_size, window_size, 1
mask_windows
=
mask_windows
.
view
(
-
1
,
self
.
window_size
*
self
.
window_size
)
mask_windows
=
mask_windows
.
view
(
-
1
,
self
.
window_size
*
self
.
window_size
)
attn_mask
=
mask_windows
.
unsqueeze
(
1
)
-
mask_windows
.
unsqueeze
(
2
)
attn_mask
=
mask_windows
.
unsqueeze
(
1
)
-
mask_windows
.
unsqueeze
(
2
)
attn_mask
=
attn_mask
.
masked_fill
(
attn_mask
!=
0
,
float
(
-
100.0
))
.
masked_fill
(
attn_mask
==
0
,
float
(
0.0
))
attn_mask
=
attn_mask
.
masked_fill
(
attn_mask
!=
0
,
float
(
-
100.0
))
.
masked_fill
(
attn_mask
==
0
,
float
(
0.0
))
return
attn_mask
return
attn_mask
def
forward
(
self
,
x
,
x_size
):
def
forward
(
self
,
x
,
x_size
):
H
,
W
=
x_size
H
,
W
=
x_size
B
,
L
,
C
=
x
.
shape
B
,
L
,
C
=
x
.
shape
# assert L == H * W, "input feature has wrong size"
# assert L == H * W, "input feature has wrong size"
shortcut
=
x
shortcut
=
x
x
=
self
.
norm1
(
x
)
x
=
self
.
norm1
(
x
)
x
=
x
.
view
(
B
,
H
,
W
,
C
)
x
=
x
.
view
(
B
,
H
,
W
,
C
)
# cyclic shift
# cyclic shift
if
self
.
shift_size
>
0
:
if
self
.
shift_size
>
0
:
shifted_x
=
torch
.
roll
(
x
,
shifts
=
(
-
self
.
shift_size
,
-
self
.
shift_size
),
dims
=
(
1
,
2
))
shifted_x
=
torch
.
roll
(
x
,
shifts
=
(
-
self
.
shift_size
,
-
self
.
shift_size
),
dims
=
(
1
,
2
))
else
:
else
:
shifted_x
=
x
shifted_x
=
x
# partition windows
# partition windows
x_windows
=
window_partition
(
shifted_x
,
self
.
window_size
)
# nW*B, window_size, window_size, C
x_windows
=
window_partition
(
shifted_x
,
self
.
window_size
)
# nW*B, window_size, window_size, C
x_windows
=
x_windows
.
view
(
-
1
,
self
.
window_size
*
self
.
window_size
,
C
)
# nW*B, window_size*window_size, C
x_windows
=
x_windows
.
view
(
-
1
,
self
.
window_size
*
self
.
window_size
,
C
)
# nW*B, window_size*window_size, C
# W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
# W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size
if
self
.
input_resolution
==
x_size
:
if
self
.
input_resolution
==
x_size
:
attn_windows
=
self
.
attn
(
x_windows
,
mask
=
self
.
attn_mask
)
# nW*B, window_size*window_size, C
attn_windows
=
self
.
attn
(
x_windows
,
mask
=
self
.
attn_mask
)
# nW*B, window_size*window_size, C
else
:
else
:
attn_windows
=
self
.
attn
(
x_windows
,
mask
=
self
.
calculate_mask
(
x_size
)
.
to
(
x
.
device
))
attn_windows
=
self
.
attn
(
x_windows
,
mask
=
self
.
calculate_mask
(
x_size
)
.
to
(
x
.
device
))
# merge windows
# merge windows
attn_windows
=
attn_windows
.
view
(
-
1
,
self
.
window_size
,
self
.
window_size
,
C
)
attn_windows
=
attn_windows
.
view
(
-
1
,
self
.
window_size
,
self
.
window_size
,
C
)
shifted_x
=
window_reverse
(
attn_windows
,
self
.
window_size
,
H
,
W
)
# B H' W' C
shifted_x
=
window_reverse
(
attn_windows
,
self
.
window_size
,
H
,
W
)
# B H' W' C
# reverse cyclic shift
# reverse cyclic shift
if
self
.
shift_size
>
0
:
if
self
.
shift_size
>
0
:
x
=
torch
.
roll
(
shifted_x
,
shifts
=
(
self
.
shift_size
,
self
.
shift_size
),
dims
=
(
1
,
2
))
x
=
torch
.
roll
(
shifted_x
,
shifts
=
(
self
.
shift_size
,
self
.
shift_size
),
dims
=
(
1
,
2
))
else
:
else
:
x
=
shifted_x
x
=
shifted_x
x
=
x
.
view
(
B
,
H
*
W
,
C
)
x
=
x
.
view
(
B
,
H
*
W
,
C
)
# FFN
# FFN
x
=
shortcut
+
self
.
drop_path
(
x
)
x
=
shortcut
+
self
.
drop_path
(
x
)
x
=
x
+
self
.
drop_path
(
self
.
mlp
(
self
.
norm2
(
x
)))
x
=
x
+
self
.
drop_path
(
self
.
mlp
(
self
.
norm2
(
x
)))
return
x
return
x
def
extra_repr
(
self
)
->
str
:
def
extra_repr
(
self
)
->
str
:
return
f
"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
\
return
f
"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
\
f
"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
f
"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}"
def
flops
(
self
):
def
flops
(
self
):
flops
=
0
flops
=
0
H
,
W
=
self
.
input_resolution
H
,
W
=
self
.
input_resolution
# norm1
# norm1
flops
+=
self
.
dim
*
H
*
W
flops
+=
self
.
dim
*
H
*
W
# W-MSA/SW-MSA
# W-MSA/SW-MSA
nW
=
H
*
W
/
self
.
window_size
/
self
.
window_size
nW
=
H
*
W
/
self
.
window_size
/
self
.
window_size
flops
+=
nW
*
self
.
attn
.
flops
(
self
.
window_size
*
self
.
window_size
)
flops
+=
nW
*
self
.
attn
.
flops
(
self
.
window_size
*
self
.
window_size
)
# mlp
# mlp
flops
+=
2
*
H
*
W
*
self
.
dim
*
self
.
dim
*
self
.
mlp_ratio
flops
+=
2
*
H
*
W
*
self
.
dim
*
self
.
dim
*
self
.
mlp_ratio
# norm2
# norm2
flops
+=
self
.
dim
*
H
*
W
flops
+=
self
.
dim
*
H
*
W
return
flops
return
flops
class
PatchMerging
(
nn
.
Module
):
class
PatchMerging
(
nn
.
Module
):
r""" Patch Merging Layer.
r""" Patch Merging Layer.
Args:
Args:
input_resolution (tuple[int]): Resolution of input feature.
input_resolution (tuple[int]): Resolution of input feature.
dim (int): Number of input channels.
dim (int): Number of input channels.
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
"""
"""
def
__init__
(
self
,
input_resolution
,
dim
,
norm_layer
=
nn
.
LayerNorm
):
def
__init__
(
self
,
input_resolution
,
dim
,
norm_layer
=
nn
.
LayerNorm
):
super
()
.
__init__
()
super
()
.
__init__
()
self
.
input_resolution
=
input_resolution
self
.
input_resolution
=
input_resolution
self
.
dim
=
dim
self
.
dim
=
dim
self
.
reduction
=
nn
.
Linear
(
4
*
dim
,
2
*
dim
,
bias
=
False
)
self
.
reduction
=
nn
.
Linear
(
4
*
dim
,
2
*
dim
,
bias
=
False
)
self
.
norm
=
norm_layer
(
4
*
dim
)
self
.
norm
=
norm_layer
(
4
*
dim
)
def
forward
(
self
,
x
):
def
forward
(
self
,
x
):
"""
"""
x: B, H*W, C
x: B, H*W, C
"""
"""
H
,
W
=
self
.
input_resolution
H
,
W
=
self
.
input_resolution
B
,
L
,
C
=
x
.
shape
B
,
L
,
C
=
x
.
shape
assert
L
==
H
*
W
,
"input feature has wrong size"
assert
L
==
H
*
W
,
"input feature has wrong size"
assert
H
%
2
==
0
and
W
%
2
==
0
,
f
"x size ({H}*{W}) are not even."
assert
H
%
2
==
0
and
W
%
2
==
0
,
f
"x size ({H}*{W}) are not even."
x
=
x
.
view
(
B
,
H
,
W
,
C
)
x
=
x
.
view
(
B
,
H
,
W
,
C
)
x0
=
x
[:,
0
::
2
,
0
::
2
,
:]
# B H/2 W/2 C
x0
=
x
[:,
0
::
2
,
0
::
2
,
:]
# B H/2 W/2 C
x1
=
x
[:,
1
::
2
,
0
::
2
,
:]
# B H/2 W/2 C
x1
=
x
[:,
1
::
2
,
0
::
2
,
:]
# B H/2 W/2 C
x2
=
x
[:,
0
::
2
,
1
::
2
,
:]
# B H/2 W/2 C
x2
=
x
[:,
0
::
2
,
1
::
2
,
:]
# B H/2 W/2 C
x3
=
x
[:,
1
::
2
,
1
::
2
,
:]
# B H/2 W/2 C
x3
=
x
[:,
1
::
2
,
1
::
2
,
:]
# B H/2 W/2 C
x
=
torch
.
cat
([
x0
,
x1
,
x2
,
x3
],
-
1
)
# B H/2 W/2 4*C
x
=
torch
.
cat
([
x0
,
x1
,
x2
,
x3
],
-
1
)
# B H/2 W/2 4*C
x
=
x
.
view
(
B
,
-
1
,
4
*
C
)
# B H/2*W/2 4*C
x
=
x
.
view
(
B
,
-
1
,
4
*
C
)
# B H/2*W/2 4*C
x
=
self
.
norm
(
x
)
x
=
self
.
norm
(
x
)
x
=
self
.
reduction
(
x
)
x
=
self
.
reduction
(
x
)
return
x
return
x
def
extra_repr
(
self
)
->
str
:
def
extra_repr
(
self
)
->
str
:
return
f
"input_resolution={self.input_resolution}, dim={self.dim}"
return
f
"input_resolution={self.input_resolution}, dim={self.dim}"
def
flops
(
self
):
def
flops
(
self
):
H
,
W
=
self
.
input_resolution
H
,
W
=
self
.
input_resolution
flops
=
H
*
W
*
self
.
dim
flops
=
H
*
W
*
self
.
dim
flops
+=
(
H
//
2
)
*
(
W
//
2
)
*
4
*
self
.
dim
*
2
*
self
.
dim
flops
+=
(
H
//
2
)
*
(
W
//
2
)
*
4
*
self
.
dim
*
2
*
self
.
dim
return
flops
return
flops
class
BasicLayer
(
nn
.
Module
):
class
BasicLayer
(
nn
.
Module
):
""" A basic Swin Transformer layer for one stage.
""" A basic Swin Transformer layer for one stage.
Args:
Args:
dim (int): Number of input channels.
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resolution.
input_resolution (tuple[int]): Input resolution.
depth (int): Number of blocks.
depth (int): Number of blocks.
num_heads (int): Number of attention heads.
num_heads (int): Number of attention heads.
window_size (int): Local window size.
window_size (int): Local window size.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
"""
"""
def
__init__
(
self
,
dim
,
input_resolution
,
depth
,
num_heads
,
window_size
,
def
__init__
(
self
,
dim
,
input_resolution
,
depth
,
num_heads
,
window_size
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop
=
0.
,
attn_drop
=
0.
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop
=
0.
,
attn_drop
=
0.
,
drop_path
=
0.
,
norm_layer
=
nn
.
LayerNorm
,
downsample
=
None
,
use_checkpoint
=
False
):
drop_path
=
0.
,
norm_layer
=
nn
.
LayerNorm
,
downsample
=
None
,
use_checkpoint
=
False
):
super
()
.
__init__
()
super
()
.
__init__
()
self
.
dim
=
dim
self
.
dim
=
dim
self
.
input_resolution
=
input_resolution
self
.
input_resolution
=
input_resolution
self
.
depth
=
depth
self
.
depth
=
depth
self
.
use_checkpoint
=
use_checkpoint
self
.
use_checkpoint
=
use_checkpoint
# build blocks
# build blocks
self
.
blocks
=
nn
.
ModuleList
([
self
.
blocks
=
nn
.
ModuleList
([
SwinTransformerBlock
(
dim
=
dim
,
input_resolution
=
input_resolution
,
SwinTransformerBlock
(
dim
=
dim
,
input_resolution
=
input_resolution
,
num_heads
=
num_heads
,
window_size
=
window_size
,
num_heads
=
num_heads
,
window_size
=
window_size
,
shift_size
=
0
if
(
i
%
2
==
0
)
else
window_size
//
2
,
shift_size
=
0
if
(
i
%
2
==
0
)
else
window_size
//
2
,
mlp_ratio
=
mlp_ratio
,
mlp_ratio
=
mlp_ratio
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
drop
=
drop
,
attn_drop
=
attn_drop
,
drop
=
drop
,
attn_drop
=
attn_drop
,
drop_path
=
drop_path
[
i
]
if
isinstance
(
drop_path
,
list
)
else
drop_path
,
drop_path
=
drop_path
[
i
]
if
isinstance
(
drop_path
,
list
)
else
drop_path
,
norm_layer
=
norm_layer
)
norm_layer
=
norm_layer
)
for
i
in
range
(
depth
)])
for
i
in
range
(
depth
)])
# patch merging layer
# patch merging layer
if
downsample
is
not
None
:
if
downsample
is
not
None
:
self
.
downsample
=
downsample
(
input_resolution
,
dim
=
dim
,
norm_layer
=
norm_layer
)
self
.
downsample
=
downsample
(
input_resolution
,
dim
=
dim
,
norm_layer
=
norm_layer
)
else
:
else
:
self
.
downsample
=
None
self
.
downsample
=
None
def
forward
(
self
,
x
,
x_size
):
def
forward
(
self
,
x
,
x_size
):
for
blk
in
self
.
blocks
:
for
blk
in
self
.
blocks
:
if
self
.
use_checkpoint
:
if
self
.
use_checkpoint
:
x
=
checkpoint
.
checkpoint
(
blk
,
x
,
x_size
)
x
=
checkpoint
.
checkpoint
(
blk
,
x
,
x_size
)
else
:
else
:
x
=
blk
(
x
,
x_size
)
x
=
blk
(
x
,
x_size
)
if
self
.
downsample
is
not
None
:
if
self
.
downsample
is
not
None
:
x
=
self
.
downsample
(
x
)
x
=
self
.
downsample
(
x
)
return
x
return
x
def
extra_repr
(
self
)
->
str
:
def
extra_repr
(
self
)
->
str
:
return
f
"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
return
f
"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
def
flops
(
self
):
def
flops
(
self
):
flops
=
0
flops
=
0
for
blk
in
self
.
blocks
:
for
blk
in
self
.
blocks
:
flops
+=
blk
.
flops
()
flops
+=
blk
.
flops
()
if
self
.
downsample
is
not
None
:
if
self
.
downsample
is
not
None
:
flops
+=
self
.
downsample
.
flops
()
flops
+=
self
.
downsample
.
flops
()
return
flops
return
flops
class
RSTB
(
nn
.
Module
):
class
RSTB
(
nn
.
Module
):
"""Residual Swin Transformer Block (RSTB).
"""Residual Swin Transformer Block (RSTB).
Args:
Args:
dim (int): Number of input channels.
dim (int): Number of input channels.
input_resolution (tuple[int]): Input resolution.
input_resolution (tuple[int]): Input resolution.
depth (int): Number of blocks.
depth (int): Number of blocks.
num_heads (int): Number of attention heads.
num_heads (int): Number of attention heads.
window_size (int): Local window size.
window_size (int): Local window size.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.
drop (float, optional): Dropout rate. Default: 0.0
drop (float, optional): Dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
attn_drop (float, optional): Attention dropout rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.
img_size: Input image size.
img_size: Input image size.
patch_size: Patch size.
patch_size: Patch size.
resi_connection: The convolutional block before residual connection.
resi_connection: The convolutional block before residual connection.
"""
"""
def
__init__
(
self
,
dim
,
input_resolution
,
depth
,
num_heads
,
window_size
,
def
__init__
(
self
,
dim
,
input_resolution
,
depth
,
num_heads
,
window_size
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop
=
0.
,
attn_drop
=
0.
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop
=
0.
,
attn_drop
=
0.
,
drop_path
=
0.
,
norm_layer
=
nn
.
LayerNorm
,
downsample
=
None
,
use_checkpoint
=
False
,
drop_path
=
0.
,
norm_layer
=
nn
.
LayerNorm
,
downsample
=
None
,
use_checkpoint
=
False
,
img_size
=
224
,
patch_size
=
4
,
resi_connection
=
'1conv'
):
img_size
=
224
,
patch_size
=
4
,
resi_connection
=
'1conv'
):
super
(
RSTB
,
self
)
.
__init__
()
super
(
RSTB
,
self
)
.
__init__
()
self
.
dim
=
dim
self
.
dim
=
dim
self
.
input_resolution
=
input_resolution
self
.
input_resolution
=
input_resolution
self
.
residual_group
=
BasicLayer
(
dim
=
dim
,
self
.
residual_group
=
BasicLayer
(
dim
=
dim
,
input_resolution
=
input_resolution
,
input_resolution
=
input_resolution
,
depth
=
depth
,
depth
=
depth
,
num_heads
=
num_heads
,
num_heads
=
num_heads
,
window_size
=
window_size
,
window_size
=
window_size
,
mlp_ratio
=
mlp_ratio
,
mlp_ratio
=
mlp_ratio
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
drop
=
drop
,
attn_drop
=
attn_drop
,
drop
=
drop
,
attn_drop
=
attn_drop
,
drop_path
=
drop_path
,
drop_path
=
drop_path
,
norm_layer
=
norm_layer
,
norm_layer
=
norm_layer
,
downsample
=
downsample
,
downsample
=
downsample
,
use_checkpoint
=
use_checkpoint
)
use_checkpoint
=
use_checkpoint
)
if
resi_connection
==
'1conv'
:
if
resi_connection
==
'1conv'
:
self
.
conv
=
nn
.
Conv2d
(
dim
,
dim
,
3
,
1
,
1
)
self
.
conv
=
nn
.
Conv2d
(
dim
,
dim
,
3
,
1
,
1
)
elif
resi_connection
==
'3conv'
:
elif
resi_connection
==
'3conv'
:
# to save parameters and memory
# to save parameters and memory
self
.
conv
=
nn
.
Sequential
(
nn
.
Conv2d
(
dim
,
dim
//
4
,
3
,
1
,
1
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
self
.
conv
=
nn
.
Sequential
(
nn
.
Conv2d
(
dim
,
dim
//
4
,
3
,
1
,
1
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
Conv2d
(
dim
//
4
,
dim
//
4
,
1
,
1
,
0
),
nn
.
Conv2d
(
dim
//
4
,
dim
//
4
,
1
,
1
,
0
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
Conv2d
(
dim
//
4
,
dim
,
3
,
1
,
1
))
nn
.
Conv2d
(
dim
//
4
,
dim
,
3
,
1
,
1
))
self
.
patch_embed
=
PatchEmbed
(
self
.
patch_embed
=
PatchEmbed
(
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
0
,
embed_dim
=
dim
,
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
0
,
embed_dim
=
dim
,
norm_layer
=
None
)
norm_layer
=
None
)
self
.
patch_unembed
=
PatchUnEmbed
(
self
.
patch_unembed
=
PatchUnEmbed
(
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
0
,
embed_dim
=
dim
,
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
0
,
embed_dim
=
dim
,
norm_layer
=
None
)
norm_layer
=
None
)
def
forward
(
self
,
x
,
x_size
):
def
forward
(
self
,
x
,
x_size
):
return
self
.
patch_embed
(
self
.
conv
(
self
.
patch_unembed
(
self
.
residual_group
(
x
,
x_size
),
x_size
)))
+
x
return
self
.
patch_embed
(
self
.
conv
(
self
.
patch_unembed
(
self
.
residual_group
(
x
,
x_size
),
x_size
)))
+
x
def
flops
(
self
):
def
flops
(
self
):
flops
=
0
flops
=
0
flops
+=
self
.
residual_group
.
flops
()
flops
+=
self
.
residual_group
.
flops
()
H
,
W
=
self
.
input_resolution
H
,
W
=
self
.
input_resolution
flops
+=
H
*
W
*
self
.
dim
*
self
.
dim
*
9
flops
+=
H
*
W
*
self
.
dim
*
self
.
dim
*
9
flops
+=
self
.
patch_embed
.
flops
()
flops
+=
self
.
patch_embed
.
flops
()
flops
+=
self
.
patch_unembed
.
flops
()
flops
+=
self
.
patch_unembed
.
flops
()
return
flops
return
flops
class
PatchEmbed
(
nn
.
Module
):
class
PatchEmbed
(
nn
.
Module
):
r""" Image to Patch Embedding
r""" Image to Patch Embedding
Args:
Args:
img_size (int): Image size. Default: 224.
img_size (int): Image size. Default: 224.
patch_size (int): Patch token size. Default: 4.
patch_size (int): Patch token size. Default: 4.
in_chans (int): Number of input image channels. Default: 3.
in_chans (int): Number of input image channels. Default: 3.
embed_dim (int): Number of linear projection output channels. Default: 96.
embed_dim (int): Number of linear projection output channels. Default: 96.
norm_layer (nn.Module, optional): Normalization layer. Default: None
norm_layer (nn.Module, optional): Normalization layer. Default: None
"""
"""
def
__init__
(
self
,
img_size
=
224
,
patch_size
=
4
,
in_chans
=
3
,
embed_dim
=
96
,
norm_layer
=
None
):
def
__init__
(
self
,
img_size
=
224
,
patch_size
=
4
,
in_chans
=
3
,
embed_dim
=
96
,
norm_layer
=
None
):
super
()
.
__init__
()
super
()
.
__init__
()
img_size
=
to_2tuple
(
img_size
)
img_size
=
to_2tuple
(
img_size
)
patch_size
=
to_2tuple
(
patch_size
)
patch_size
=
to_2tuple
(
patch_size
)
patches_resolution
=
[
img_size
[
0
]
//
patch_size
[
0
],
img_size
[
1
]
//
patch_size
[
1
]]
patches_resolution
=
[
img_size
[
0
]
//
patch_size
[
0
],
img_size
[
1
]
//
patch_size
[
1
]]
self
.
img_size
=
img_size
self
.
img_size
=
img_size
self
.
patch_size
=
patch_size
self
.
patch_size
=
patch_size
self
.
patches_resolution
=
patches_resolution
self
.
patches_resolution
=
patches_resolution
self
.
num_patches
=
patches_resolution
[
0
]
*
patches_resolution
[
1
]
self
.
num_patches
=
patches_resolution
[
0
]
*
patches_resolution
[
1
]
self
.
in_chans
=
in_chans
self
.
in_chans
=
in_chans
self
.
embed_dim
=
embed_dim
self
.
embed_dim
=
embed_dim
if
norm_layer
is
not
None
:
if
norm_layer
is
not
None
:
self
.
norm
=
norm_layer
(
embed_dim
)
self
.
norm
=
norm_layer
(
embed_dim
)
else
:
else
:
self
.
norm
=
None
self
.
norm
=
None
def
forward
(
self
,
x
):
def
forward
(
self
,
x
):
x
=
x
.
flatten
(
2
)
.
transpose
(
1
,
2
)
# B Ph*Pw C
x
=
x
.
flatten
(
2
)
.
transpose
(
1
,
2
)
# B Ph*Pw C
if
self
.
norm
is
not
None
:
if
self
.
norm
is
not
None
:
x
=
self
.
norm
(
x
)
x
=
self
.
norm
(
x
)
return
x
return
x
def
flops
(
self
):
def
flops
(
self
):
flops
=
0
flops
=
0
H
,
W
=
self
.
img_size
H
,
W
=
self
.
img_size
if
self
.
norm
is
not
None
:
if
self
.
norm
is
not
None
:
flops
+=
H
*
W
*
self
.
embed_dim
flops
+=
H
*
W
*
self
.
embed_dim
return
flops
return
flops
class
PatchUnEmbed
(
nn
.
Module
):
class
PatchUnEmbed
(
nn
.
Module
):
r""" Image to Patch Unembedding
r""" Image to Patch Unembedding
Args:
Args:
img_size (int): Image size. Default: 224.
img_size (int): Image size. Default: 224.
patch_size (int): Patch token size. Default: 4.
patch_size (int): Patch token size. Default: 4.
in_chans (int): Number of input image channels. Default: 3.
in_chans (int): Number of input image channels. Default: 3.
embed_dim (int): Number of linear projection output channels. Default: 96.
embed_dim (int): Number of linear projection output channels. Default: 96.
norm_layer (nn.Module, optional): Normalization layer. Default: None
norm_layer (nn.Module, optional): Normalization layer. Default: None
"""
"""
def
__init__
(
self
,
img_size
=
224
,
patch_size
=
4
,
in_chans
=
3
,
embed_dim
=
96
,
norm_layer
=
None
):
def
__init__
(
self
,
img_size
=
224
,
patch_size
=
4
,
in_chans
=
3
,
embed_dim
=
96
,
norm_layer
=
None
):
super
()
.
__init__
()
super
()
.
__init__
()
img_size
=
to_2tuple
(
img_size
)
img_size
=
to_2tuple
(
img_size
)
patch_size
=
to_2tuple
(
patch_size
)
patch_size
=
to_2tuple
(
patch_size
)
patches_resolution
=
[
img_size
[
0
]
//
patch_size
[
0
],
img_size
[
1
]
//
patch_size
[
1
]]
patches_resolution
=
[
img_size
[
0
]
//
patch_size
[
0
],
img_size
[
1
]
//
patch_size
[
1
]]
self
.
img_size
=
img_size
self
.
img_size
=
img_size
self
.
patch_size
=
patch_size
self
.
patch_size
=
patch_size
self
.
patches_resolution
=
patches_resolution
self
.
patches_resolution
=
patches_resolution
self
.
num_patches
=
patches_resolution
[
0
]
*
patches_resolution
[
1
]
self
.
num_patches
=
patches_resolution
[
0
]
*
patches_resolution
[
1
]
self
.
in_chans
=
in_chans
self
.
in_chans
=
in_chans
self
.
embed_dim
=
embed_dim
self
.
embed_dim
=
embed_dim
def
forward
(
self
,
x
,
x_size
):
def
forward
(
self
,
x
,
x_size
):
B
,
HW
,
C
=
x
.
shape
B
,
HW
,
C
=
x
.
shape
x
=
x
.
transpose
(
1
,
2
)
.
view
(
B
,
self
.
embed_dim
,
x_size
[
0
],
x_size
[
1
])
# B Ph*Pw C
x
=
x
.
transpose
(
1
,
2
)
.
view
(
B
,
self
.
embed_dim
,
x_size
[
0
],
x_size
[
1
])
# B Ph*Pw C
return
x
return
x
def
flops
(
self
):
def
flops
(
self
):
flops
=
0
flops
=
0
return
flops
return
flops
class
Upsample
(
nn
.
Sequential
):
class
Upsample
(
nn
.
Sequential
):
"""Upsample module.
"""Upsample module.
Args:
Args:
scale (int): Scale factor. Supported scales: 2^n and 3.
scale (int): Scale factor. Supported scales: 2^n and 3.
num_feat (int): Channel number of intermediate features.
num_feat (int): Channel number of intermediate features.
"""
"""
def
__init__
(
self
,
scale
,
num_feat
):
def
__init__
(
self
,
scale
,
num_feat
):
m
=
[]
m
=
[]
if
(
scale
&
(
scale
-
1
))
==
0
:
# scale = 2^n
if
(
scale
&
(
scale
-
1
))
==
0
:
# scale = 2^n
for
_
in
range
(
int
(
math
.
log
(
scale
,
2
))):
for
_
in
range
(
int
(
math
.
log
(
scale
,
2
))):
m
.
append
(
nn
.
Conv2d
(
num_feat
,
4
*
num_feat
,
3
,
1
,
1
))
m
.
append
(
nn
.
Conv2d
(
num_feat
,
4
*
num_feat
,
3
,
1
,
1
))
m
.
append
(
nn
.
PixelShuffle
(
2
))
m
.
append
(
nn
.
PixelShuffle
(
2
))
elif
scale
==
3
:
elif
scale
==
3
:
m
.
append
(
nn
.
Conv2d
(
num_feat
,
9
*
num_feat
,
3
,
1
,
1
))
m
.
append
(
nn
.
Conv2d
(
num_feat
,
9
*
num_feat
,
3
,
1
,
1
))
m
.
append
(
nn
.
PixelShuffle
(
3
))
m
.
append
(
nn
.
PixelShuffle
(
3
))
else
:
else
:
raise
ValueError
(
f
'scale {scale} is not supported. '
'Supported scales: 2^n and 3.'
)
raise
ValueError
(
f
'scale {scale} is not supported. '
'Supported scales: 2^n and 3.'
)
super
(
Upsample
,
self
)
.
__init__
(
*
m
)
super
(
Upsample
,
self
)
.
__init__
(
*
m
)
class
UpsampleOneStep
(
nn
.
Sequential
):
class
UpsampleOneStep
(
nn
.
Sequential
):
"""UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
"""UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)
Used in lightweight SR to save parameters.
Used in lightweight SR to save parameters.
Args:
Args:
scale (int): Scale factor. Supported scales: 2^n and 3.
scale (int): Scale factor. Supported scales: 2^n and 3.
num_feat (int): Channel number of intermediate features.
num_feat (int): Channel number of intermediate features.
"""
"""
def
__init__
(
self
,
scale
,
num_feat
,
num_out_ch
,
input_resolution
=
None
):
def
__init__
(
self
,
scale
,
num_feat
,
num_out_ch
,
input_resolution
=
None
):
self
.
num_feat
=
num_feat
self
.
num_feat
=
num_feat
self
.
input_resolution
=
input_resolution
self
.
input_resolution
=
input_resolution
m
=
[]
m
=
[]
m
.
append
(
nn
.
Conv2d
(
num_feat
,
(
scale
**
2
)
*
num_out_ch
,
3
,
1
,
1
))
m
.
append
(
nn
.
Conv2d
(
num_feat
,
(
scale
**
2
)
*
num_out_ch
,
3
,
1
,
1
))
m
.
append
(
nn
.
PixelShuffle
(
scale
))
m
.
append
(
nn
.
PixelShuffle
(
scale
))
super
(
UpsampleOneStep
,
self
)
.
__init__
(
*
m
)
super
(
UpsampleOneStep
,
self
)
.
__init__
(
*
m
)
def
flops
(
self
):
def
flops
(
self
):
H
,
W
=
self
.
input_resolution
H
,
W
=
self
.
input_resolution
flops
=
H
*
W
*
self
.
num_feat
*
3
*
9
flops
=
H
*
W
*
self
.
num_feat
*
3
*
9
return
flops
return
flops
class
SwinIR
(
nn
.
Module
):
class
SwinIR
(
nn
.
Module
):
r""" SwinIR
r""" SwinIR
A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.
A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.
Args:
Args:
img_size (int | tuple(int)): Input image size. Default 64
img_size (int | tuple(int)): Input image size. Default 64
patch_size (int | tuple(int)): Patch size. Default: 1
patch_size (int | tuple(int)): Patch size. Default: 1
in_chans (int): Number of input image channels. Default: 3
in_chans (int): Number of input image channels. Default: 3
embed_dim (int): Patch embedding dimension. Default: 96
embed_dim (int): Patch embedding dimension. Default: 96
depths (tuple(int)): Depth of each Swin Transformer layer.
depths (tuple(int)): Depth of each Swin Transformer layer.
num_heads (tuple(int)): Number of attention heads in different layers.
num_heads (tuple(int)): Number of attention heads in different layers.
window_size (int): Window size. Default: 7
window_size (int): Window size. Default: 7
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4
qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True
qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None
drop_rate (float): Dropout rate. Default: 0
drop_rate (float): Dropout rate. Default: 0
attn_drop_rate (float): Attention dropout rate. Default: 0
attn_drop_rate (float): Attention dropout rate. Default: 0
drop_path_rate (float): Stochastic depth rate. Default: 0.1
drop_path_rate (float): Stochastic depth rate. Default: 0.1
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.
ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
ape (bool): If True, add absolute position embedding to the patch embedding. Default: False
patch_norm (bool): If True, add normalization after patch embedding. Default: True
patch_norm (bool): If True, add normalization after patch embedding. Default: True
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False
upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction
img_range: Image range. 1. or 255.
img_range: Image range. 1. or 255.
upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None
resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
resi_connection: The convolutional block before residual connection. '1conv'/'3conv'
"""
"""
def
__init__
(
self
,
img_size
=
64
,
patch_size
=
1
,
in_chans
=
3
,
def
__init__
(
self
,
img_size
=
64
,
patch_size
=
1
,
in_chans
=
3
,
embed_dim
=
96
,
depths
=
[
6
,
6
,
6
,
6
],
num_heads
=
[
6
,
6
,
6
,
6
],
embed_dim
=
96
,
depths
=
[
6
,
6
,
6
,
6
],
num_heads
=
[
6
,
6
,
6
,
6
],
window_size
=
7
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
window_size
=
7
,
mlp_ratio
=
4.
,
qkv_bias
=
True
,
qk_scale
=
None
,
drop_rate
=
0.
,
attn_drop_rate
=
0.
,
drop_path_rate
=
0.1
,
drop_rate
=
0.
,
attn_drop_rate
=
0.
,
drop_path_rate
=
0.1
,
norm_layer
=
nn
.
LayerNorm
,
ape
=
False
,
patch_norm
=
True
,
norm_layer
=
nn
.
LayerNorm
,
ape
=
False
,
patch_norm
=
True
,
use_checkpoint
=
False
,
upscale
=
2
,
img_range
=
1.
,
upsampler
=
''
,
resi_connection
=
'1conv'
,
use_checkpoint
=
False
,
upscale
=
2
,
img_range
=
1.
,
upsampler
=
''
,
resi_connection
=
'1conv'
,
**
kwargs
):
**
kwargs
):
super
(
SwinIR
,
self
)
.
__init__
()
super
(
SwinIR
,
self
)
.
__init__
()
num_in_ch
=
in_chans
num_in_ch
=
in_chans
num_out_ch
=
in_chans
num_out_ch
=
in_chans
num_feat
=
64
num_feat
=
64
self
.
img_range
=
img_range
self
.
img_range
=
img_range
if
in_chans
==
3
:
if
in_chans
==
3
:
rgb_mean
=
(
0.4488
,
0.4371
,
0.4040
)
rgb_mean
=
(
0.4488
,
0.4371
,
0.4040
)
self
.
mean
=
torch
.
Tensor
(
rgb_mean
)
.
view
(
1
,
3
,
1
,
1
)
self
.
mean
=
torch
.
Tensor
(
rgb_mean
)
.
view
(
1
,
3
,
1
,
1
)
else
:
else
:
self
.
mean
=
torch
.
zeros
(
1
,
1
,
1
,
1
)
self
.
mean
=
torch
.
zeros
(
1
,
1
,
1
,
1
)
self
.
upscale
=
upscale
self
.
upscale
=
upscale
self
.
upsampler
=
upsampler
self
.
upsampler
=
upsampler
self
.
window_size
=
window_size
self
.
window_size
=
window_size
#####################################################################################################
#####################################################################################################
################################### 1, shallow feature extraction ###################################
################################### 1, shallow feature extraction ###################################
self
.
conv_first
=
nn
.
Conv2d
(
num_in_ch
,
embed_dim
,
3
,
1
,
1
)
self
.
conv_first
=
nn
.
Conv2d
(
num_in_ch
,
embed_dim
,
3
,
1
,
1
)
#####################################################################################################
#####################################################################################################
################################### 2, deep feature extraction ######################################
################################### 2, deep feature extraction ######################################
self
.
num_layers
=
len
(
depths
)
self
.
num_layers
=
len
(
depths
)
self
.
embed_dim
=
embed_dim
self
.
embed_dim
=
embed_dim
self
.
ape
=
ape
self
.
ape
=
ape
self
.
patch_norm
=
patch_norm
self
.
patch_norm
=
patch_norm
self
.
num_features
=
embed_dim
self
.
num_features
=
embed_dim
self
.
mlp_ratio
=
mlp_ratio
self
.
mlp_ratio
=
mlp_ratio
# split image into non-overlapping patches
# split image into non-overlapping patches
self
.
patch_embed
=
PatchEmbed
(
self
.
patch_embed
=
PatchEmbed
(
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
embed_dim
,
embed_dim
=
embed_dim
,
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
embed_dim
,
embed_dim
=
embed_dim
,
norm_layer
=
norm_layer
if
self
.
patch_norm
else
None
)
norm_layer
=
norm_layer
if
self
.
patch_norm
else
None
)
num_patches
=
self
.
patch_embed
.
num_patches
num_patches
=
self
.
patch_embed
.
num_patches
patches_resolution
=
self
.
patch_embed
.
patches_resolution
patches_resolution
=
self
.
patch_embed
.
patches_resolution
self
.
patches_resolution
=
patches_resolution
self
.
patches_resolution
=
patches_resolution
# merge non-overlapping patches into image
# merge non-overlapping patches into image
self
.
patch_unembed
=
PatchUnEmbed
(
self
.
patch_unembed
=
PatchUnEmbed
(
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
embed_dim
,
embed_dim
=
embed_dim
,
img_size
=
img_size
,
patch_size
=
patch_size
,
in_chans
=
embed_dim
,
embed_dim
=
embed_dim
,
norm_layer
=
norm_layer
if
self
.
patch_norm
else
None
)
norm_layer
=
norm_layer
if
self
.
patch_norm
else
None
)
# absolute position embedding
# absolute position embedding
if
self
.
ape
:
if
self
.
ape
:
self
.
absolute_pos_embed
=
nn
.
Parameter
(
torch
.
zeros
(
1
,
num_patches
,
embed_dim
))
self
.
absolute_pos_embed
=
nn
.
Parameter
(
torch
.
zeros
(
1
,
num_patches
,
embed_dim
))
trunc_normal_
(
self
.
absolute_pos_embed
,
std
=
.02
)
trunc_normal_
(
self
.
absolute_pos_embed
,
std
=
.02
)
self
.
pos_drop
=
nn
.
Dropout
(
p
=
drop_rate
)
self
.
pos_drop
=
nn
.
Dropout
(
p
=
drop_rate
)
# stochastic depth
# stochastic depth
dpr
=
[
x
.
item
()
for
x
in
torch
.
linspace
(
0
,
drop_path_rate
,
sum
(
depths
))]
# stochastic depth decay rule
dpr
=
[
x
.
item
()
for
x
in
torch
.
linspace
(
0
,
drop_path_rate
,
sum
(
depths
))]
# stochastic depth decay rule
# build Residual Swin Transformer blocks (RSTB)
# build Residual Swin Transformer blocks (RSTB)
self
.
layers
=
nn
.
ModuleList
()
self
.
layers
=
nn
.
ModuleList
()
for
i_layer
in
range
(
self
.
num_layers
):
for
i_layer
in
range
(
self
.
num_layers
):
layer
=
RSTB
(
dim
=
embed_dim
,
layer
=
RSTB
(
dim
=
embed_dim
,
input_resolution
=
(
patches_resolution
[
0
],
input_resolution
=
(
patches_resolution
[
0
],
patches_resolution
[
1
]),
patches_resolution
[
1
]),
depth
=
depths
[
i_layer
],
depth
=
depths
[
i_layer
],
num_heads
=
num_heads
[
i_layer
],
num_heads
=
num_heads
[
i_layer
],
window_size
=
window_size
,
window_size
=
window_size
,
mlp_ratio
=
self
.
mlp_ratio
,
mlp_ratio
=
self
.
mlp_ratio
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
qkv_bias
=
qkv_bias
,
qk_scale
=
qk_scale
,
drop
=
drop_rate
,
attn_drop
=
attn_drop_rate
,
drop
=
drop_rate
,
attn_drop
=
attn_drop_rate
,
drop_path
=
dpr
[
sum
(
depths
[:
i_layer
]):
sum
(
depths
[:
i_layer
+
1
])],
# no impact on SR results
drop_path
=
dpr
[
sum
(
depths
[:
i_layer
]):
sum
(
depths
[:
i_layer
+
1
])],
# no impact on SR results
norm_layer
=
norm_layer
,
norm_layer
=
norm_layer
,
downsample
=
None
,
downsample
=
None
,
use_checkpoint
=
use_checkpoint
,
use_checkpoint
=
use_checkpoint
,
img_size
=
img_size
,
img_size
=
img_size
,
patch_size
=
patch_size
,
patch_size
=
patch_size
,
resi_connection
=
resi_connection
resi_connection
=
resi_connection
)
)
self
.
layers
.
append
(
layer
)
self
.
layers
.
append
(
layer
)
self
.
norm
=
norm_layer
(
self
.
num_features
)
self
.
norm
=
norm_layer
(
self
.
num_features
)
# build the last conv layer in deep feature extraction
# build the last conv layer in deep feature extraction
if
resi_connection
==
'1conv'
:
if
resi_connection
==
'1conv'
:
self
.
conv_after_body
=
nn
.
Conv2d
(
embed_dim
,
embed_dim
,
3
,
1
,
1
)
self
.
conv_after_body
=
nn
.
Conv2d
(
embed_dim
,
embed_dim
,
3
,
1
,
1
)
elif
resi_connection
==
'3conv'
:
elif
resi_connection
==
'3conv'
:
# to save parameters and memory
# to save parameters and memory
self
.
conv_after_body
=
nn
.
Sequential
(
nn
.
Conv2d
(
embed_dim
,
embed_dim
//
4
,
3
,
1
,
1
),
self
.
conv_after_body
=
nn
.
Sequential
(
nn
.
Conv2d
(
embed_dim
,
embed_dim
//
4
,
3
,
1
,
1
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
Conv2d
(
embed_dim
//
4
,
embed_dim
//
4
,
1
,
1
,
0
),
nn
.
Conv2d
(
embed_dim
//
4
,
embed_dim
//
4
,
1
,
1
,
0
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
),
nn
.
Conv2d
(
embed_dim
//
4
,
embed_dim
,
3
,
1
,
1
))
nn
.
Conv2d
(
embed_dim
//
4
,
embed_dim
,
3
,
1
,
1
))
#####################################################################################################
#####################################################################################################
################################ 3, high quality image reconstruction ################################
################################ 3, high quality image reconstruction ################################
if
self
.
upsampler
==
'pixelshuffle'
:
if
self
.
upsampler
==
'pixelshuffle'
:
# for classical SR
# for classical SR
self
.
conv_before_upsample
=
nn
.
Sequential
(
nn
.
Conv2d
(
embed_dim
,
num_feat
,
3
,
1
,
1
),
self
.
conv_before_upsample
=
nn
.
Sequential
(
nn
.
Conv2d
(
embed_dim
,
num_feat
,
3
,
1
,
1
),
nn
.
LeakyReLU
(
inplace
=
True
))
nn
.
LeakyReLU
(
inplace
=
True
))
self
.
upsample
=
Upsample
(
upscale
,
num_feat
)
self
.
upsample
=
Upsample
(
upscale
,
num_feat
)
self
.
conv_last
=
nn
.
Conv2d
(
num_feat
,
num_out_ch
,
3
,
1
,
1
)
self
.
conv_last
=
nn
.
Conv2d
(
num_feat
,
num_out_ch
,
3
,
1
,
1
)
elif
self
.
upsampler
==
'pixelshuffledirect'
:
elif
self
.
upsampler
==
'pixelshuffledirect'
:
# for lightweight SR (to save parameters)
# for lightweight SR (to save parameters)
self
.
upsample
=
UpsampleOneStep
(
upscale
,
embed_dim
,
num_out_ch
,
self
.
upsample
=
UpsampleOneStep
(
upscale
,
embed_dim
,
num_out_ch
,
(
patches_resolution
[
0
],
patches_resolution
[
1
]))
(
patches_resolution
[
0
],
patches_resolution
[
1
]))
elif
self
.
upsampler
==
'nearest+conv'
:
elif
self
.
upsampler
==
'nearest+conv'
:
# for real-world SR (less artifacts)
# for real-world SR (less artifacts)
self
.
conv_before_upsample
=
nn
.
Sequential
(
nn
.
Conv2d
(
embed_dim
,
num_feat
,
3
,
1
,
1
),
self
.
conv_before_upsample
=
nn
.
Sequential
(
nn
.
Conv2d
(
embed_dim
,
num_feat
,
3
,
1
,
1
),
nn
.
LeakyReLU
(
inplace
=
True
))
nn
.
LeakyReLU
(
inplace
=
True
))
self
.
conv_up1
=
nn
.
Conv2d
(
num_feat
,
num_feat
,
3
,
1
,
1
)
self
.
conv_up1
=
nn
.
Conv2d
(
num_feat
,
num_feat
,
3
,
1
,
1
)
if
self
.
upscale
==
4
:
if
self
.
upscale
==
4
:
self
.
conv_up2
=
nn
.
Conv2d
(
num_feat
,
num_feat
,
3
,
1
,
1
)
self
.
conv_up2
=
nn
.
Conv2d
(
num_feat
,
num_feat
,
3
,
1
,
1
)
self
.
conv_hr
=
nn
.
Conv2d
(
num_feat
,
num_feat
,
3
,
1
,
1
)
self
.
conv_hr
=
nn
.
Conv2d
(
num_feat
,
num_feat
,
3
,
1
,
1
)
self
.
conv_last
=
nn
.
Conv2d
(
num_feat
,
num_out_ch
,
3
,
1
,
1
)
self
.
conv_last
=
nn
.
Conv2d
(
num_feat
,
num_out_ch
,
3
,
1
,
1
)
self
.
lrelu
=
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
)
self
.
lrelu
=
nn
.
LeakyReLU
(
negative_slope
=
0.2
,
inplace
=
True
)
else
:
else
:
# for image denoising and JPEG compression artifact reduction
# for image denoising and JPEG compression artifact reduction
self
.
conv_last
=
nn
.
Conv2d
(
embed_dim
,
num_out_ch
,
3
,
1
,
1
)
self
.
conv_last
=
nn
.
Conv2d
(
embed_dim
,
num_out_ch
,
3
,
1
,
1
)
self
.
apply
(
self
.
_init_weights
)
self
.
apply
(
self
.
_init_weights
)
def
_init_weights
(
self
,
m
):
def
_init_weights
(
self
,
m
):
if
isinstance
(
m
,
nn
.
Linear
):
if
isinstance
(
m
,
nn
.
Linear
):
trunc_normal_
(
m
.
weight
,
std
=
.02
)
trunc_normal_
(
m
.
weight
,
std
=
.02
)
if
isinstance
(
m
,
nn
.
Linear
)
and
m
.
bias
is
not
None
:
if
isinstance
(
m
,
nn
.
Linear
)
and
m
.
bias
is
not
None
:
nn
.
init
.
constant_
(
m
.
bias
,
0
)
nn
.
init
.
constant_
(
m
.
bias
,
0
)
elif
isinstance
(
m
,
nn
.
LayerNorm
):
elif
isinstance
(
m
,
nn
.
LayerNorm
):
nn
.
init
.
constant_
(
m
.
bias
,
0
)
nn
.
init
.
constant_
(
m
.
bias
,
0
)
nn
.
init
.
constant_
(
m
.
weight
,
1.0
)
nn
.
init
.
constant_
(
m
.
weight
,
1.0
)
@
torch
.
jit
.
ignore
@
torch
.
jit
.
ignore
def
no_weight_decay
(
self
):
def
no_weight_decay
(
self
):
return
{
'absolute_pos_embed'
}
return
{
'absolute_pos_embed'
}
@
torch
.
jit
.
ignore
@
torch
.
jit
.
ignore
def
no_weight_decay_keywords
(
self
):
def
no_weight_decay_keywords
(
self
):
return
{
'relative_position_bias_table'
}
return
{
'relative_position_bias_table'
}
def
check_image_size
(
self
,
x
):
def
check_image_size
(
self
,
x
):
_
,
_
,
h
,
w
=
x
.
size
()
_
,
_
,
h
,
w
=
x
.
size
()
mod_pad_h
=
(
self
.
window_size
-
h
%
self
.
window_size
)
%
self
.
window_size
mod_pad_h
=
(
self
.
window_size
-
h
%
self
.
window_size
)
%
self
.
window_size
mod_pad_w
=
(
self
.
window_size
-
w
%
self
.
window_size
)
%
self
.
window_size
mod_pad_w
=
(
self
.
window_size
-
w
%
self
.
window_size
)
%
self
.
window_size
x
=
F
.
pad
(
x
,
(
0
,
mod_pad_w
,
0
,
mod_pad_h
),
'reflect'
)
x
=
F
.
pad
(
x
,
(
0
,
mod_pad_w
,
0
,
mod_pad_h
),
'reflect'
)
return
x
return
x
def
forward_features
(
self
,
x
):
def
forward_features
(
self
,
x
):
x_size
=
(
x
.
shape
[
2
],
x
.
shape
[
3
])
x_size
=
(
x
.
shape
[
2
],
x
.
shape
[
3
])
x
=
self
.
patch_embed
(
x
)
x
=
self
.
patch_embed
(
x
)
if
self
.
ape
:
if
self
.
ape
:
x
=
x
+
self
.
absolute_pos_embed
x
=
x
+
self
.
absolute_pos_embed
x
=
self
.
pos_drop
(
x
)
x
=
self
.
pos_drop
(
x
)
for
layer
in
self
.
layers
:
for
layer
in
self
.
layers
:
x
=
layer
(
x
,
x_size
)
x
=
layer
(
x
,
x_size
)
x
=
self
.
norm
(
x
)
# B L C
x
=
self
.
norm
(
x
)
# B L C
x
=
self
.
patch_unembed
(
x
,
x_size
)
x
=
self
.
patch_unembed
(
x
,
x_size
)
return
x
return
x
def
forward
(
self
,
x
):
def
forward
(
self
,
x
):
H
,
W
=
x
.
shape
[
2
:]
H
,
W
=
x
.
shape
[
2
:]
x
=
self
.
check_image_size
(
x
)
x
=
self
.
check_image_size
(
x
)
self
.
mean
=
self
.
mean
.
type_as
(
x
)
self
.
mean
=
self
.
mean
.
type_as
(
x
)
x
=
(
x
-
self
.
mean
)
*
self
.
img_range
x
=
(
x
-
self
.
mean
)
*
self
.
img_range
if
self
.
upsampler
==
'pixelshuffle'
:
if
self
.
upsampler
==
'pixelshuffle'
:
# for classical SR
# for classical SR
x
=
self
.
conv_first
(
x
)
x
=
self
.
conv_first
(
x
)
x
=
self
.
conv_after_body
(
self
.
forward_features
(
x
))
+
x
x
=
self
.
conv_after_body
(
self
.
forward_features
(
x
))
+
x
x
=
self
.
conv_before_upsample
(
x
)
x
=
self
.
conv_before_upsample
(
x
)
x
=
self
.
conv_last
(
self
.
upsample
(
x
))
x
=
self
.
conv_last
(
self
.
upsample
(
x
))
elif
self
.
upsampler
==
'pixelshuffledirect'
:
elif
self
.
upsampler
==
'pixelshuffledirect'
:
# for lightweight SR
# for lightweight SR
x
=
self
.
conv_first
(
x
)
x
=
self
.
conv_first
(
x
)
x
=
self
.
conv_after_body
(
self
.
forward_features
(
x
))
+
x
x
=
self
.
conv_after_body
(
self
.
forward_features
(
x
))
+
x
x
=
self
.
upsample
(
x
)
x
=
self
.
upsample
(
x
)
elif
self
.
upsampler
==
'nearest+conv'
:
elif
self
.
upsampler
==
'nearest+conv'
:
# for real-world SR
# for real-world SR
x
=
self
.
conv_first
(
x
)
x
=
self
.
conv_first
(
x
)
x
=
self
.
conv_after_body
(
self
.
forward_features
(
x
))
+
x
x
=
self
.
conv_after_body
(
self
.
forward_features
(
x
))
+
x
x
=
self
.
conv_before_upsample
(
x
)
x
=
self
.
conv_before_upsample
(
x
)
x
=
self
.
lrelu
(
self
.
conv_up1
(
torch
.
nn
.
functional
.
interpolate
(
x
,
scale_factor
=
2
,
mode
=
'nearest'
)))
x
=
self
.
lrelu
(
self
.
conv_up1
(
torch
.
nn
.
functional
.
interpolate
(
x
,
scale_factor
=
2
,
mode
=
'nearest'
)))
if
self
.
upscale
==
4
:
if
self
.
upscale
==
4
:
x
=
self
.
lrelu
(
self
.
conv_up2
(
torch
.
nn
.
functional
.
interpolate
(
x
,
scale_factor
=
2
,
mode
=
'nearest'
)))
x
=
self
.
lrelu
(
self
.
conv_up2
(
torch
.
nn
.
functional
.
interpolate
(
x
,
scale_factor
=
2
,
mode
=
'nearest'
)))
x
=
self
.
conv_last
(
self
.
lrelu
(
self
.
conv_hr
(
x
)))
x
=
self
.
conv_last
(
self
.
lrelu
(
self
.
conv_hr
(
x
)))
else
:
else
:
# for image denoising and JPEG compression artifact reduction
# for image denoising and JPEG compression artifact reduction
x_first
=
self
.
conv_first
(
x
)
x_first
=
self
.
conv_first
(
x
)
res
=
self
.
conv_after_body
(
self
.
forward_features
(
x_first
))
+
x_first
res
=
self
.
conv_after_body
(
self
.
forward_features
(
x_first
))
+
x_first
x
=
x
+
self
.
conv_last
(
res
)
x
=
x
+
self
.
conv_last
(
res
)
x
=
x
/
self
.
img_range
+
self
.
mean
x
=
x
/
self
.
img_range
+
self
.
mean
return
x
[:,
:,
:
H
*
self
.
upscale
,
:
W
*
self
.
upscale
]
return
x
[:,
:,
:
H
*
self
.
upscale
,
:
W
*
self
.
upscale
]
def
flops
(
self
):
def
flops
(
self
):
flops
=
0
flops
=
0
H
,
W
=
self
.
patches_resolution
H
,
W
=
self
.
patches_resolution
flops
+=
H
*
W
*
3
*
self
.
embed_dim
*
9
flops
+=
H
*
W
*
3
*
self
.
embed_dim
*
9
flops
+=
self
.
patch_embed
.
flops
()
flops
+=
self
.
patch_embed
.
flops
()
for
i
,
layer
in
enumerate
(
self
.
layers
):
for
i
,
layer
in
enumerate
(
self
.
layers
):
flops
+=
layer
.
flops
()
flops
+=
layer
.
flops
()
flops
+=
H
*
W
*
3
*
self
.
embed_dim
*
self
.
embed_dim
flops
+=
H
*
W
*
3
*
self
.
embed_dim
*
self
.
embed_dim
flops
+=
self
.
upsample
.
flops
()
flops
+=
self
.
upsample
.
flops
()
return
flops
return
flops
if
__name__
==
'__main__'
:
if
__name__
==
'__main__'
:
upscale
=
4
upscale
=
4
window_size
=
8
window_size
=
8
height
=
(
1024
//
upscale
//
window_size
+
1
)
*
window_size
height
=
(
1024
//
upscale
//
window_size
+
1
)
*
window_size
width
=
(
720
//
upscale
//
window_size
+
1
)
*
window_size
width
=
(
720
//
upscale
//
window_size
+
1
)
*
window_size
model
=
SwinIR
(
upscale
=
2
,
img_size
=
(
height
,
width
),
model
=
SwinIR
(
upscale
=
2
,
img_size
=
(
height
,
width
),
window_size
=
window_size
,
img_range
=
1.
,
depths
=
[
6
,
6
,
6
,
6
],
window_size
=
window_size
,
img_range
=
1.
,
depths
=
[
6
,
6
,
6
,
6
],
embed_dim
=
60
,
num_heads
=
[
6
,
6
,
6
,
6
],
mlp_ratio
=
2
,
upsampler
=
'pixelshuffledirect'
)
embed_dim
=
60
,
num_heads
=
[
6
,
6
,
6
,
6
],
mlp_ratio
=
2
,
upsampler
=
'pixelshuffledirect'
)
print
(
model
)
print
(
model
)
print
(
height
,
width
,
model
.
flops
()
/
1e9
)
print
(
height
,
width
,
model
.
flops
()
/
1e9
)
x
=
torch
.
randn
((
1
,
3
,
height
,
width
))
x
=
torch
.
randn
((
1
,
3
,
height
,
width
))
x
=
model
(
x
)
x
=
model
(
x
)
print
(
x
.
shape
)
print
(
x
.
shape
)
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment