flotingdream / trainner-redux-fj Goto Github PK

(Fork of BasicSR joeyballentine/traiNNer-redux) Open Source Image and Video Restoration Toolbox for Super-resolution, Denoise, Deblurring, etc. Currently, it includes EDSR, RCAN, SRResNet, SRGAN, ESRGAN, EDVR, BasicVSR, SwinIR, ECBSR, OmniSR, HAT, GRL, A-ESRGAN, DAT, WaveMixSR, StarSRGAN, DLGSANet, DITN etc. Also support StyleGAN2, DFDNet.

NEW ADD ARCH SUPPORT

• [ESRGAN 8X]
• Omnisr
  • The arch implementation of Omnisr is from Omnisr. The LICENSE of Omnisr is [Apache License 2.0]. The LICENSE is included as LICENSE_Omnisr.
• HAT
  • The arch implementation of HAT is from HAT. The LICENSE of HAT is [MIT License]. The LICENSE is included as LICENSE_HAT.
• GRL
  • The arch implementation of GRL is from GRL. The LICENSE of GRL is [MIT License].
• ESWT # have bug
  • The arch implementation of ESWT is from ESWT. The LICENSE of ESWT is [MIT License].
• SRFormer
  • The arch implementation of SRFormer is from SRFormer. The LICENSE of SRFormer is [Apache License 2.0].
• A-ESRGAN
  • The arch implementation of A-ESRGAN is from A-ESRGAN. The LICENSE of A-ESRGAN is [BSD 3-Clause "New" or "Revised" License].
• DAT
  • The arch implementation of DAT is from DAT. The LICENSE of DAT is [Apache License 2.0].
• WaveMixSR
  • The arch implementation of WaveMixSR is from WaveMixSR. The LICENSE of WaveMixSR is [MIT License].
• StarSRGAN
  • The arch implementation of StarSRGAN is from StarSRGAN. The LICENSE of StarSRGAN is [Apache License 2.0].
• DLGSANet
  • The arch implementation of DLGSANet is from DLGSANet. The LICENSE of DLGSANet is [Apache License 2.0].
• DITN
  • The arch implementation of DITN is from DITN.

NEW FEATURE SUPPORT

• ContextualLoss weight
• amp support

  try_autoamp_g: True # enable amp Automatic mixed precision for network_g. if loss inf or nan or error just set to False
  try_autoamp_d: True # enable amp Automatic mixed precision for network_d. if loss inf or nan or error just set to False
  

TIPS

• PairedImageDataset set high_order_degradation : False
• RealESRGANDataset set high_order_degradation : True
• model_type can use GeneralGANModel or GeneralNetModel for ,ost of archs
• To use Automatic mixed precision, edit the yml
• e.g. if you want use omnisr-gan with amp, just use model_type: GeneralGANModel and edit the .yml 's following part

name: train_OmniSRGANModel_SRx4_scratch_P48W8_DIV2K_500k_B4G8
try_autoamp_g: True # enable amp Automatic mixed precision for network_g. if loss inf or nan or error just set to False
try_autoamp_d: True # enable amp Automatic mixed precision for network_d. if loss inf or nan or error just set to False
network_g:
  type: OmniSRNet
  num_in_ch: 3
  num_out_ch: 3
  num_feat: 64
  upsampling: 4
  res_num: 5
  block_num: 1
  bias: True
  block_script_name: OSA
  block_class_name: OSA_Block
  window_size: 8
  pe: True
  ffn_bias: True

• If you want to use aesrgan's network_d for other network_d. you should edit the .yml network_d to multiscale type
• e.g. if you want use omnisr as network_g and multiscale as network_d, just edit the .yml 's following part

# network structures

network_g:
  type: OmniSRNet
  num_in_ch: 3
  num_out_ch: 3
  num_feat: 64
  upsampling: 4
  res_num: 5
  block_num: 1
  bias: True
  block_script_name: OSA
  block_class_name: OSA_Block
  window_size: 8
  pe: True
  ffn_bias: True

network_d:
  type: multiscale
  num_in_ch: 3
  num_feat: 64
  num_D: 2

for easy use here are examples for network_g

# network structures
network_g:
  # DAT-S, need to set batch size >1
  type: DAT
  upscale: 4
  in_chans: 3
  img_size: 64
  img_range: 1.
  split_size: [8,16]
  depth: [6,6,6,6,6,6]
  embed_dim: 180
  num_heads: [6,6,6,6,6,6]
  expansion_factor: 2
  resi_connection: '1conv'

  type: WaveMixSR
  scale: 4
  depth: 4
  mult: 1
  ff_channel: 144
  final_dim: 144
  dropout: 0.3

  type: StarSRNet
  num_in_ch: 3
  num_out_ch: 3
  scale: 4
  num_feat: 64
  num_block: 23
  num_grow_ch: 32
  drop_out: False

  type: DLGSANet
  upscale: 4
  in_chans: 3
  dim: 90
  groups: 3
  blocks: 2
  buildblock_type: 'sparseedge'
  window_size: 7
  idynamic_num_heads: 6
  idynamic_ffn_type: 'GDFN'
  idynamic_ffn_expansion_factor: 2.
  idynamic: true
  restormer_num_heads: 6
  restormer_ffn_type: 'GDFN'
  restormer_ffn_expansion_factor: 2.
  tlc_flag: true
  tlc_kernel: 48    # using tlc during validation
  activation: 'relu'
  body_norm: false
  img_range: 1.
  upsampler: 'pixelshuffledirect'

  type: DITN # need torch >2.0
  inp_channels: 3
  dim:  60
  ITL_blocks:  4
  SAL_blocks:  4
  UFONE_blocks:  1
  ffn_expansion_factor:  2
  bias:  False
  LayerNorm_type:  'WithBias'
  patch_size: 8
  upscale: 4