convert tf batch normalization to pytorch

1import numpy as np
2import torch
3import torch.nn as nn
4from torch.autograd import Variable
5
6class Model(nn.Module):
7    def __init__(self):
8        super(Model, self).__init__()
9        self.module_list = nn.ModuleList()
10        module = nn.Sequential()
11
12        conv = nn.Conv2d(3, 32, 3, 1, 1, bias=False)
13        module.add_module('conv_0', conv)
14
15        bn = nn.BatchNorm2d(32)
16        module.add_module('batch_norm_0', bn)
17
18        gamma = np.random.rand(32)
19        gamma = torch.from_numpy(gamma)
20        bn.weight.data.copy_(gamma)
21
22        beta = np.random.rand(32)
23        beta = torch.from_numpy(beta)
24        bn.bias.data.copy_(beta)
25
26        mean = np.random.rand(32)
27        mean = torch.from_numpy(mean)
28        bn.running_mean.data.copy_(mean)
29
30        var = np.random.rand(32)
31        var = torch.from_numpy(var)
32        bn.running_var.data.copy_(var)
33
34        self.module_list.append(module)
35
36    def forward(self, input):
37        conv = self.module_list[0][0](input)
38        bn = self.module_list[0][1](conv)
39        return conv, bn
40
41
42if __name__ == '__main__':
43 	x = np.random.rand(1, 3, 64, 64)
44    x = Variable(torch.from_numpy(x).float())
45    
46    model = Model()    
47    model.eval()
48    
49    with torch.no_grad():
50		conv_out, bn_out = model.forward(x)
51
52