"docs/source/vscode:/vscode.git/clone" did not exist on "9a7235faf2835d424c4587e703024248e6b9f465"
Commit cbb05c59 authored by Marat Dukhan's avatar Marat Dukhan Committed by Francisco Massa
Browse files

Use torch.nn.init in SqueezeNet models (#146)

parent 874481f2
import math import math
import torch import torch
import torch.nn as nn import torch.nn as nn
import torch.nn.init as init
import torch.utils.model_zoo as model_zoo import torch.utils.model_zoo as model_zoo
...@@ -87,13 +88,10 @@ class SqueezeNet(nn.Module): ...@@ -87,13 +88,10 @@ class SqueezeNet(nn.Module):
for m in self.modules(): for m in self.modules():
if isinstance(m, nn.Conv2d): if isinstance(m, nn.Conv2d):
gain = 2.0
if m is final_conv: if m is final_conv:
m.weight.data.normal_(0, 0.01) init.normal(m.weight.data, mean=0.0, std=0.01)
else: else:
fan_in = m.kernel_size[0] * m.kernel_size[1] * m.in_channels init.kaiming_uniform(m.weight.data)
u = math.sqrt(3.0 * gain / fan_in)
m.weight.data.uniform_(-u, u)
if m.bias is not None: if m.bias is not None:
m.bias.data.zero_() m.bias.data.zero_()
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment