aboutsummaryrefslogtreecommitdiff
path: root/extensions-builtin/SwinIR/swinir_model_arch.py
diff options
context:
space:
mode:
authorAUTOMATIC1111 <16777216c@gmail.com>2023-05-22 07:15:34 +0300
committerGitHub <noreply@github.com>2023-05-22 07:15:34 +0300
commit8137bdba61fd57cc1ddae801f6080d51e13d70c5 (patch)
treec5a02e9f9ae57c9f0ff8499379c6cc61a97c094e /extensions-builtin/SwinIR/swinir_model_arch.py
parenta862428902c4aecde8852761c3a4d95c196885cb (diff)
parent3366e494a1147e570d8527eea19da88edb3a1e0c (diff)
Merge branch 'dev' into text-drag-fix
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()