-
Notifications
You must be signed in to change notification settings - Fork 1
/
resnet18.py
73 lines (57 loc) · 1.99 KB
/
resnet18.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
66
67
68
69
70
71
72
73
import torch
from torch import nn as nn
from src.modules.base_generator import GeneratorAbstract
from torchvision import models
class Resnet18(nn.Module):
def __init__(self, in_channel: int, out_channel: int, pretrained: bool):
"""
Args:
in_channel: input channels.
out_channel: output channels.
"""
super().__init__()
self.out_channel = out_channel
self.model = models.resnet18(pretrained=pretrained)
del self.model.fc
del self.model.avgpool
if self.out_channel == 512:
pass
elif self.out_channel == 256:
del self.model.layer4
elif self.out_channel == 128:
del self.model.layer4
del self.model.layer3
elif self.out_channel == 64:
del self.model.layer4
del self.model.layer3
del self.model.layer2
else:
raise Exception("out_channel: 512, 256, 128 or 64")
def forward(self,x):
x = self.model.conv1(x)
x = self.model.bn1(x)
x = self.model.relu(x)
x = self.model.maxpool(x)
x = self.model.layer1(x)
if self.out_channel >= 128:
x = self.model.layer2(x)
if self.out_channel >= 256:
x = self.model.layer3(x)
if self.out_channel >= 512:
x = self.model.layer4(x)
return x
class Resnet18Generator(GeneratorAbstract):
""" Resnet18 (torchvision.models) module generator for parsing."""
def __init__(self, *args, **kwargs):
"""Initailize."""
super().__init__(*args, **kwargs)
@property
def out_channel(self) -> int:
"""Get out channel size."""
return self.args[0]
def __call__(self, repeat: int = 1):
# TODO: Apply repeat
pretrained = self.args[1] if len(self.args) > 1 else True
return self._get_module(
Resnet18(self.in_channel, self.out_channel, pretrained=pretrained)
)