-
Notifications
You must be signed in to change notification settings - Fork 10
/
BaseCNN.py
65 lines (57 loc) · 2.14 KB
/
BaseCNN.py
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
import torch.nn as nn
from torchvision import models
from BCNN import BCNN
class BaseCNN(nn.Module):
def __init__(self, config):
"""Declare all needed layers."""
nn.Module.__init__(self)
self.config = config
if self.config.backbone == 'resnet18':
self.backbone = models.resnet18(pretrained=True)
elif self.config.backbone == 'resnet34':
self.backbone = models.resnet34(pretrained=True)
# elif self.config.backbone == 'resnet50':
# self.backbone = models.resnet50(pretrained=True)
# self.fc = nn.Linear(2048, 1)
if config.std_modeling:
outdim = 2
else:
outdim = 1
if config.representation == 'BCNN':
assert ((self.config.backbone == 'resnet18') | (self.config.backbone == 'resnet34')), "The backbone network must be resnet18 or resnet34"
self.representation = BCNN()
self.fc = nn.Linear(512 * 512, outdim)
else:
self.fc = nn.Linear(512, outdim)
if self.config.fc:
# Freeze all previous layers.
for param in self.backbone.parameters():
param.requires_grad = False
# Initialize the fc layers.
nn.init.kaiming_normal_(self.fc.weight.data)
if self.fc.bias is not None:
nn.init.constant_(self.fc.bias.data, val=0)
def forward(self, x):
"""Forward pass of the network.
"""
x = self.backbone.conv1(x)
x = self.backbone.bn1(x)
x = self.backbone.relu(x)
x = self.backbone.maxpool(x)
x = self.backbone.layer1(x)
x = self.backbone.layer2(x)
x = self.backbone.layer3(x)
x = self.backbone.layer4(x)
if self.config.representation == 'BCNN':
x = self.representation(x)
else:
x = self.backbone.avgpool(x)
x = torch.flatten(x, start_dim=1)
x = self.fc(x)
if self.config.std_modeling:
mean = x[:, 0]
t = x[:, 1]
var = nn.functional.softplus(t)
return mean, var
else:
return x