suggestions and fixes from the PR

This commit is contained in:
AUTOMATIC
2023-05-10 21:21:32 +03:00
parent d25219b7e8
commit 3ec7b705c7
11 changed files with 16 additions and 31 deletions

View File

@@ -644,17 +644,13 @@ class SwinIR(nn.Module):
"""
def __init__(self, img_size=64, patch_size=1, in_chans=3,
embed_dim=96, depths=None, num_heads=None,
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,
use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',
**kwargs):
super(SwinIR, self).__init__()
depths = depths or [6, 6, 6, 6]
num_heads = num_heads or [6, 6, 6, 6]
num_in_ch = in_chans
num_out_ch = in_chans
num_feat = 64