How to initialize weights in PyTorch?

Question

How to initialize the weights and biases (for example, with He or Xavier initialization) in a network in PyTorch?

Answer 1

If you see a deprecation warning (@Fábio Perez)...

def init_weights(m):
    if type(m) == nn.Linear:
        torch.nn.init.xavier_uniform_(m.weight)
        m.bias.data.fill_(0.01)

net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
net.apply(init_weights)