基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)

前面完成了基于知识蒸馏的去雨去雪去雾模型大的部署与训练,下面则进行代码的学习。
使用debug的方式进行代码的学习。
首先是网络结构展示:轻易不要打开,这个模型太复杂了。说到底倒不是多复杂,就是层数太多了

Net(
  (conv_input): ConvLayer(
    (reflection_pad): ReflectionPad2d((5, 5, 5, 5))
    (conv2d): Conv2d(3, 16, kernel_size=(11, 11), stride=(1, 1))
  )
  (dense0): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv2x): ConvLayer(
    (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
    (conv2d): Conv2d(16, 32, kernel_size=(3, 3), stride=(2, 2))
  )
  (conv1): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(48, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(64, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(80, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion1): Encoder_MDCBlock1(
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(16, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (dense1): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv4x): ConvLayer(
    (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
    (conv2d): Conv2d(32, 64, kernel_size=(3, 3), stride=(2, 2))
  )
  (conv2): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(160, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion2): Encoder_MDCBlock1(
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): DeconvBlock(
        (deconv): ConvTranspose2d(16, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): ConvBlock(
        (conv): Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (dense2): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv8x): ConvLayer(
    (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
    (conv2d): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2))
  )
  (conv3): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(256, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion3): Encoder_MDCBlock1(
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): DeconvBlock(
        (deconv): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): DeconvBlock(
        (deconv): ConvTranspose2d(16, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): ConvBlock(
        (conv): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): ConvBlock(
        (conv): Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (dense3): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv16x): ConvLayer(
    (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
    (conv2d): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2))
  )
  (conv4): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(384, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(512, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(640, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion4): Encoder_MDCBlock1(
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): DeconvBlock(
        (deconv): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): DeconvBlock(
        (deconv): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (3): DeconvBlock(
        (deconv): ConvTranspose2d(16, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): ConvBlock(
        (conv): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): ConvBlock(
        (conv): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (3): ConvBlock(
        (conv): Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (dehaze): Sequential(
    (res0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res3): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res4): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res5): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res6): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res7): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res8): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res9): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res10): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res11): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res12): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res13): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res14): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res15): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res16): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (res17): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (convd16x): UpsampleConvLayer(
    (conv2d): ConvTranspose2d(256, 128, kernel_size=(3, 3), stride=(2, 2))
  )
  (dense_4): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv_4): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(192, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(256, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(320, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion_4): Decoder_MDCBlock1(
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (convd8x): UpsampleConvLayer(
    (conv2d): ConvTranspose2d(128, 64, kernel_size=(3, 3), stride=(2, 2))
  )
  (dense_3): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv_3): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(96, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(128, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(160, 32, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion_3): Decoder_MDCBlock1(
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): ConvBlock(
        (conv): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): DeconvBlock(
        (deconv): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (convd4x): UpsampleConvLayer(
    (conv2d): ConvTranspose2d(64, 32, kernel_size=(3, 3), stride=(2, 2))
  )
  (dense_2): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv_2): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(32, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(48, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(64, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(80, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion_2): Decoder_MDCBlock1(
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): ConvBlock(
        (conv): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): ConvBlock(
        (conv): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): DeconvBlock(
        (deconv): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): DeconvBlock(
        (deconv): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (convd2x): UpsampleConvLayer(
    (conv2d): ConvTranspose2d(32, 16, kernel_size=(3, 3), stride=(2, 2))
  )
  (dense_1): Sequential(
    (0): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (1): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
    (2): ResidualBlock(
      (conv1): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (conv2): ConvLayer(
        (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
        (conv2d): Conv2d(16, 16, kernel_size=(3, 3), stride=(1, 1))
      )
      (relu): PReLU(num_parameters=1)
    )
  )
  (conv_1): RDB(
    (dense_layers): Sequential(
      (0): make_dense(
        (conv): Conv2d(8, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (1): make_dense(
        (conv): Conv2d(16, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (2): make_dense(
        (conv): Conv2d(24, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
      (3): make_dense(
        (conv): Conv2d(32, 8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
    (conv_1x1): Conv2d(40, 8, kernel_size=(1, 1), stride=(1, 1), bias=False)
  )
  (fusion_1): Decoder_MDCBlock1(
    (down_convs): ModuleList(
      (0): ConvBlock(
        (conv): Conv2d(8, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): ConvBlock(
        (conv): Conv2d(16, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): ConvBlock(
        (conv): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (3): ConvBlock(
        (conv): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
    (up_convs): ModuleList(
      (0): DeconvBlock(
        (deconv): ConvTranspose2d(16, 8, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (1): DeconvBlock(
        (deconv): ConvTranspose2d(32, 16, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (2): DeconvBlock(
        (deconv): ConvTranspose2d(64, 32, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
      (3): DeconvBlock(
        (deconv): ConvTranspose2d(128, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
        (act): PReLU(num_parameters=1)
      )
    )
  )
  (conv_output): ConvLayer(
    (reflection_pad): ReflectionPad2d((1, 1, 1, 1))
    (conv2d): Conv2d(16, 3, kernel_size=(3, 3), stride=(1, 1))
  )
)

首先进入训练模式,又称知识收集训练阶段:

def train_kc_stage(model, teacher_networks, ckt_modules, train_loader, optimizer, scheduler, epoch, criterions):
	print(Fore.CYAN + "==> Training Stage 1")
	print("==> Epoch {}/{}".format(epoch, args.max_epoch))
	print("==> Learning Rate = {:.6f}".format(optimizer.param_groups[0]['lr']))
	meters = get_meter(num_meters=5)	
	criterion_l1, criterion_scr, _ = criterions
	model.train()
	ckt_modules.train()
	for teacher_network in teacher_networks:
		teacher_network.eval()

声明所需要的损失函数,ckt_models(协作知识迁移模型)的训练模式
ckt_models 的详细结构如下:

ModuleList(
  (0): CKTModule(
    (teacher_projectors): TeacherProjectors(
      (PFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
      (IPFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
    )
    (student_projector): StudentProjector(
      (PFP): Sequential(
        (0): Conv2d(64, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (1): ReLU(inplace=True)
        (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
  )
  (1): CKTModule(
    (teacher_projectors): TeacherProjectors(
      (PFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
      (IPFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
    )
    (student_projector): StudentProjector(
      (PFP): Sequential(
        (0): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (1): ReLU(inplace=True)
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
  )
  (2): CKTModule(
    (teacher_projectors): TeacherProjectors(
      (PFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
      (IPFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
    )
    (student_projector): StudentProjector(
      (PFP): Sequential(
        (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (1): ReLU(inplace=True)
        (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
  )
  (3): CKTModule(
    (teacher_projectors): TeacherProjectors(
      (PFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
      (IPFPs): ModuleList(
        (0): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (1): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
        (2): Sequential(
          (0): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
          (1): ReLU(inplace=True)
          (2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        )
      )
    )
    (student_projector): StudentProjector(
      (PFP): Sequential(
        (0): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
        (1): ReLU(inplace=True)
        (2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
      )
    )
  )
)

criterions的结构,其定义的是损失函数,分别是L1损失,SCR损失以及HCR损失

ModuleList(
  (0): L1Loss()
  (1): SCRLoss(
    (vgg): Vgg19(
      (slice1): Sequential(
        (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU(inplace=True)
      )
      (slice2): Sequential(
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU(inplace=True)
        (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (6): ReLU(inplace=True)
      )
      (slice3): Sequential(
        (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (8): ReLU(inplace=True)
        (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (11): ReLU(inplace=True)
      )
      (slice4): Sequential(
        (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (13): ReLU(inplace=True)
        (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (15): ReLU(inplace=True)
        (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (17): ReLU(inplace=True)
        (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (20): ReLU(inplace=True)
      )
      (slice5): Sequential(
        (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (22): ReLU(inplace=True)
        (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (24): ReLU(inplace=True)
        (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (26): ReLU(inplace=True)
        (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (29): ReLU(inplace=True)
      )
    )
    (l1): L1Loss()
  )
  (2): HCRLoss(
    (vgg): Vgg19(
      (slice1): Sequential(
        (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (1): ReLU(inplace=True)
      )
      (slice2): Sequential(
        (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (3): ReLU(inplace=True)
        (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (6): ReLU(inplace=True)
      )
      (slice3): Sequential(
        (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (8): ReLU(inplace=True)
        (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (11): ReLU(inplace=True)
      )
      (slice4): Sequential(
        (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (13): ReLU(inplace=True)
        (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (15): ReLU(inplace=True)
        (16): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (17): ReLU(inplace=True)
        (18): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (19): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (20): ReLU(inplace=True)
      )
      (slice5): Sequential(
        (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (22): ReLU(inplace=True)
        (23): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (24): ReLU(inplace=True)
        (25): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (26): ReLU(inplace=True)
        (27): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
        (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
        (29): ReLU(inplace=True)
      )
    )
    (l1): L1Loss()
  )
)

可以看到教师网络就是将先前的Net网络复制了3份,只是加载不同权重而已。即三个model。

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第1张图片

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第2张图片

继续训练

start = time.time()
pBar = tqdm(train_loader, desc='Training')
for target_images, input_images in pBar:
	
	# Check whether the batch contains all types of degraded data
	if target_images is None: continue

	# move to GPU
	target_images = target_images.cuda()
	input_images = [images.cuda() for images in input_images]

	# Fix all teachers and collect reconstruction results and features from cooresponding teacher
	preds_from_teachers = []
	features_from_each_teachers = []
	with torch.no_grad():
		for i in range(len(teacher_networks)):
			preds, features = teacher_networks[i](input_images[i], return_feat=True)
			preds_from_teachers.append(preds)
			features_from_each_teachers.append(features)	
			
	preds_from_teachers = torch.cat(preds_from_teachers)
	features_from_teachers = []
	for layer in range(len(features_from_each_teachers[0])):
		features_from_teachers.append([features_from_each_teachers[i][layer] for i in range(len(teacher_networks))])

	preds_from_student, features_from_student = model(torch.cat(input_images), return_feat=True)   

	
	# Project the features to common feature space and calculate the loss
	PFE_loss, PFV_loss = 0., 0.
	for i, (s_features, t_features) in enumerate(zip(features_from_student, features_from_teachers)):
		t_proj_features, t_recons_features, s_proj_features = ckt_modules[i](t_features, s_features)
		PFE_loss += criterion_l1(s_proj_features, torch.cat(t_proj_features))
		PFV_loss += 0.05 * criterion_l1(torch.cat(t_recons_features), torch.cat(t_features))

	T_loss = criterion_l1(preds_from_student, preds_from_teachers)
	SCR_loss = 0.1 * criterion_scr(preds_from_student, target_images, torch.cat(input_images))
	total_loss = T_loss + PFE_loss + PFV_loss + SCR_loss

	optimizer.zero_grad()
	total_loss.backward()
	optimizer.step()

进入评估模块:加载模型,验证集,最终输出psnr与ssim

if epoch % args.val_freq == 0:
			psnr, ssim = evaluate(model, val_loader, epoch)
			# Check whether the model is top-k model
			top_k_state = save_top_k(model, optimizer, scheduler, top_k_state, args.top_k, epoch, args.save_dir, psnr=psnr, ssim=ssim)

evaluate(model, val_loader, epoch) 函数详细代码:

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第3张图片
随后进行结果输出:

pred = model(image)

即跳入Net的forward中进行特征提取

输入值:

输入x: 图像维度为640x480,此时初始维度:torch.Size([1, 3, 480, 640])
随后经过一系列的卷积降维,生成了如下特征图:这个过程就不赘述了。

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第4张图片

输出值:

输出x与feature:最终的x的维度依旧为torch.Size([1, 3, 480, 640])

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第5张图片

feature的维度,共有4个特征图,分别如下:

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第6张图片

这里设置只输出x,所以pred的值即为x的值:

在这里插入图片描述
得到输出值后,即可进行损失的计算了:

psnr_list.append(torchPSNR(pred, target).item())
ssim_list.append(pytorch_ssim.ssim(pred, target).item())

具体实现:

@torch.no_grad()
def torchPSNR(prd_img, tar_img):
	if not isinstance(prd_img, torch.Tensor):
		prd_img = torch.from_numpy(prd_img)
		tar_img = torch.from_numpy(tar_img)

	imdff = torch.clamp(prd_img, 0, 1) - torch.clamp(tar_img, 0, 1)
	rmse = (imdff**2).mean().sqrt()
	ps = 20 * torch.log10(1/rmse)
	return ps

最终将19张图片全部评估完毕:

在这里插入图片描述
得到psnr_list值:
在这里插入图片描述
需要19张全部评估完,这里只进行了两张。

最终返回平均值:

return np.mean(psnr_list), np.mean(ssim_list)

该方法最终的值变为:

基于知识蒸馏的两阶段去雨去雪去雾模型学习记录(一)_第7张图片

你可能感兴趣的:(图像处理,图像去噪,学习,人工智能,python)