pixel-distribution-learning/net.py at main · youwei1-sudo/pixel-distribution-learning · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
import torch.nn.functional as F
from torch import nn
import torch


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(6, 256, 6, 2)
       # self.batch_norm1 = nn.BatchNorm2d(64)
        self.conv2 = nn.Conv2d(256, 512, 6, 2)
        # self.batch_norm2 = nn.BatchNorm2d(128)
        # self.conv3 = nn.Conv2d(512, 512, 2, 2)

        self.fc1 = nn.Linear(512, 512)
        # self.fc3 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(512, 2)

    def forward(self, x):
        # Pass data through conv1
        x = self.conv1(x)
        # x = self.batch_norm1(x)
        x = F.relu(x)
        # x = F.max_pool2d(x, kernel_size=2, stride=2)
        # Use the rectified-linear activation function over x

        x = self.conv2(x)
        # x = self.batch_norm2(x)
        x = F.relu(x)
        # x = F.max_pool2d(x, kernel_size=2, stride=2)
        # x = self.conv3(x)
        # x = F.relu(x)
        # Flatten x with start_dim=1
        x = torch.flatten(x, 1)
        # Pass data through fc1
        # try to add relu to linear layer
        x = self.fc1(x)
        # x = self.fc3(x)
        x = self.fc2(x)

        # Apply softmax to x
        output = F.log_softmax(x, dim=1)
        return output

# class Net(nn.Module):
#     def __init__(self):
#         super(Net, self).__init__()
#         self.conv1 = nn.Conv2d(6, 256, 4, 2)
#        # self.batch_norm1 = nn.BatchNorm2d(64)
#         self.conv2 = nn.Conv2d(256, 512, 4, 2)
#         # self.batch_norm2 = nn.BatchNorm2d(128)
#         self.conv3 = nn.Conv2d(512, 512, 2, 2)
#
#         self.fc1 = nn.Linear(512, 512)
#         # self.fc3 = nn.Linear(128, 64)
#         self.fc2 = nn.Linear(512, 2)
#
#     def forward(self, x):
#         # Pass data through conv1
#         x = self.conv1(x)
#         # x = self.batch_norm1(x)
#         x = F.relu(x)
#         # x = F.max_pool2d(x, kernel_size=2, stride=2)
#         # Use the rectified-linear activation function over x
#
#         x = self.conv2(x)
#         # x = self.batch_norm2(x)
#         x = F.relu(x)
#         # x = F.max_pool2d(x, kernel_size=2, stride=2)
#         x = self.conv3(x)
#         x = F.relu(x)
#         # Flatten x with start_dim=1
#         x = torch.flatten(x, 1)
#         # Pass data through fc1
#         # try to add relu to linear layer
#         x = self.fc1(x)
#         # x = self.fc3(x)
#         x = self.fc2(x)
#
#         # Apply softmax to x
#         output = F.log_softmax(x, dim=1)
#         return output

input = torch.randn(size=(1, 6, 16, 16))
net = Net()
out = net.forward(input)

pytorch_total_params = sum(p.numel() for p in net.parameters())
print(pytorch_total_params)