Search code examples
pythonmodelpytorchcaffe

What can I use in Pytorch to remplace caffe's weight filler


I'm writing a pytorch model based on the caffe model below. Do you know how I can write weight filler and bias filler in pytorch ?

layer {
  name: "conv3"
  type: "Convolution"
  bottom: "pool2"
  top: "conv3"
  param {
    lr_mult: 1
    decay_mult: 1
  }
  convolution_param {
    num_output: 128
    pad_h: 0
    pad_w: 1
    kernel_h: 1
    kernel_w: 3
    stride: 1
    weight_filler {
      type: "gaussian"
      std: 0.1
    }
    bias_filler {
      type: "constant"
      value: 0
    }
  }
}

Thank you


Solution

  • Pytorch has torch.nn.init library to help with init weights of a network.
    You probably want to use nn.init.normal_ for the "gaussian" filler, and nn.init.constant_ for the "constant" filler of the bias.

    You can use a function to fill the weights of a module m:

    def init_weights(m):
        if type(m) == nn.Conv2d:
            torch.nn.init.normal_(m.weight, std=0.1)
            if m.bias is not None:
                torch.nn.init.constant_(m.bias, val=0)
    
    # define the net
    net = MyCaffeLikeNetwork()
    # use the function to init all weights of the net
    net.apply(init_weights)
    

    For more information on weight init in pytorch you can look at this detailed answer.