aboutsummaryrefslogtreecommitdiff
path: root/extensions-builtin/SwinIR/swinir_model_arch.py
diff options
context:
space:
mode:
authorAUTOMATIC1111 <16777216c@gmail.com>2023-05-17 23:18:56 +0300
committerGitHub <noreply@github.com>2023-05-17 23:18:56 +0300
commit04b4508a66de58c9f3a422fdcad4dd2ec3ad39ce (patch)
tree1cb30a63099a69f678d4901b495203b765a6df59 /extensions-builtin/SwinIR/swinir_model_arch.py
parent7201d940a4fe664beb9662fadbeade4ee1d788f7 (diff)
parentb397f63e00bbfbe9087d80abb457aa9a593b181b (diff)
Merge branch 'dev' into improve-frontend-responsiveness
Diffstat (limited to 'extensions-builtin/SwinIR/swinir_model_arch.py')
-rw-r--r--extensions-builtin/SwinIR/swinir_model_arch.py6
1 files changed, 3 insertions, 3 deletions
diff --git a/extensions-builtin/SwinIR/swinir_model_arch.py b/extensions-builtin/SwinIR/swinir_model_arch.py
index 863f42db..93b93274 100644
--- a/extensions-builtin/SwinIR/swinir_model_arch.py
+++ b/extensions-builtin/SwinIR/swinir_model_arch.py
@@ -644,7 +644,7 @@ class SwinIR(nn.Module):
"""
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,
drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
norm_layer=nn.LayerNorm, ape=False, patch_norm=True,
@@ -805,7 +805,7 @@ class SwinIR(nn.Module):
def forward(self, x):
H, W = x.shape[2:]
x = self.check_image_size(x)
-
+
self.mean = self.mean.type_as(x)
x = (x - self.mean) * self.img_range
@@ -844,7 +844,7 @@ class SwinIR(nn.Module):
H, W = self.patches_resolution
flops += H * W * 3 * self.embed_dim * 9
flops += self.patch_embed.flops()
- for i, layer in enumerate(self.layers):
+ for layer in self.layers:
flops += layer.flops()
flops += H * W * 3 * self.embed_dim * self.embed_dim
flops += self.upsample.flops()