JingyunLiang / SwinIR

SwinIR: Image Restoration Using Swin Transformer (official repository)

Home Page:https://arxiv.org/abs/2108.10257

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Problem when saving the model

hcleung3325 opened this issue · comments

Hi thanks for the training code.
I have a problem when iteration meet 5000 to save the model.
File "/KAIR/models/network_swinir.py", line 254, in forward
x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C
File "/KAIR/models/network_swinir.py", line 42, in window_partition
x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)
RuntimeError: shape '[1, 111, 8, 143, 8, 180]' is invalid for input of size 184459500
May I know how to fix it?
Thanks.

{
  "task": "swinir_sr_classical_patch64_x4_l1"     //  classical image sr for x2/x3/x4/x8. root/task/images-models-options
  , "model": "plain" // "plain" | "plain2" if two inputs
  , "gpu_ids": [6,7]
  , "dist": true

  , "scale": 4       // 2 | 3 | 4 | 8
  , "n_channels": 3  // broadcast to "datasets", 1 for grayscale, 3 for color

  , "path": {
    "root": "superresolution"            // "denoising" | "superresolution" | "dejpeg"
    , "pretrained_netG": null      // path of pretrained model. We fine-tune X3/X4/X8 models from X2 model, so that `G_optimizer_lr` and `G_scheduler_milestones` can be halved to save time.
    , "pretrained_netE": null      // path of pretrained model
  }

  , "datasets": {
    "train": {
      "name": "train_dataset"           // just name
      , "dataset_type": "sr"         // "dncnn" | "dnpatch" | "fdncnn" | "ffdnet" | "sr" | "srmd" | "dpsr" | "plain" | "plainpatch" | "jpeg"
      , "dataroot_H": "../HR"// path of H training dataset. DIV2K (800 training images)
      , "dataroot_L": "../LR"              // path of L training dataset

      , "H_size": 256                   // 96/144|192/384 | 128/192/256/512. LR patch size is set to 48 or 64 when compared with RCAN or RRDB.

      , "dataloader_shuffle": true
      , "dataloader_num_workers": 16
      , "dataloader_batch_size": 8      // batch size 1 | 16 | 32 | 48 | 64 | 128. Total batch size =4x8=32 in SwinIR
    }
    , "test": {
      "name": "test_dataset"            // just name
      , "dataset_type": "sr"         // "dncnn" | "dnpatch" | "fdncnn" | "ffdnet" | "sr" | "srmd" | "dpsr" | "plain" | "plainpatch" | "jpeg"
      , "dataroot_H": "testsets/Set5/HR"  // path of H testing dataset
      , "dataroot_L": "testsets/Set5/LR_bicubic/X2"              // path of L testing dataset

    }
  }

  , "netG": {
    "net_type": "swinir" 
    , "upscale": 4                      // 2 | 3  | 4 | 8
    , "in_chans": 3 
    , "img_size": 64                    // For fair comparison, LR patch size is set to 48 or 64 when compared with RCAN or RRDB.
    , "window_size": 8  
    , "img_range": 1.0 
    , "depths": [6, 6, 6, 6, 6, 6] 
    , "embed_dim": 180 
    , "num_heads": [6, 6, 6, 6, 6, 6]
    , "mlp_ratio": 2 
    , "upsampler": "pixelshuffle"        // "pixelshuffle" | "pixelshuffledirect" | "nearest+conv" | null
    , "resi_connection": "1conv"        // "1conv" | "3conv"

    , "init_type": "default"
  }

  , "train": {
    "G_lossfn_type": "l1"               // "l1" preferred | "l2sum" | "l2" | "ssim" | "charbonnier"
    , "G_lossfn_weight": 1.0            // default

    , "E_decay": 0.999                  // Exponential Moving Average for netG: set 0 to disable; default setting 0.999

    , "G_optimizer_type": "adam"        // fixed, adam is enough
    , "G_optimizer_lr": 2e-4            // learning rate
    , "G_optimizer_wd": 0               // weight decay, default 0
    , "G_optimizer_clipgrad": null      // unused
    , "G_optimizer_reuse": true         // 

    , "G_scheduler_type": "MultiStepLR" // "MultiStepLR" is enough
    , "G_scheduler_milestones": [250000, 400000, 450000, 475000, 500000]
    , "G_scheduler_gamma": 0.5

    , "G_regularizer_orthstep": null    // unused
    , "G_regularizer_clipstep": null    // unused

    , "G_param_strict": true
    , "E_param_strict": true

    , "checkpoint_test": 5000           // for testing
    , "checkpoint_save": 5000           // for saving model
    , "checkpoint_print": 200           // for print
  }
}

Same problem! How to address it? Thx!

The model is "upscale": 4 , but the LR image input is "dataroot_L": "testsets/Set5/LR_bicubic/X2"