-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
65 lines (45 loc) · 1.46 KB
/
model.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
from typing import Callable, List, Tuple
import os
import torch
import catalyst
from catalyst.dl import utils
SEED = 42
utils.set_global_seed(SEED)
utils.prepare_cudnn(deterministic=True)
from torch import nn
import torch.nn.functional as F
import pretrainedmodels
from efficientnet_pytorch import EfficientNet
import timm
"""
EfficientNetB0 - (224, 224, 3)
EfficientNetB1 - (240, 240, 3)
EfficientNetB2 - (260, 260, 3)
EfficientNetB3 - (300, 300, 3)
EfficientNetB4 - (380, 380, 3)
EfficientNetB5 - (456, 456, 3)
EfficientNetB6 - (528, 528, 3)
EfficientNetB7 - (600, 600, 3)
"""
# 'mobilenetv2_100',
# 'mobilenetv2_110d',
# 'mobilenetv2_120d',
# 'mobilenetv2_140',
# 'mobilenetv3_large_100',
# 'mobilenetv3_rw',
class mobilenet(torch.nn.Module):
def __init__(self, name, output_layer = 5):
super(mobilenet, self).__init__()
self.model = timm.create_model(name, pretrained=True)
num_features = self.model.classifier.in_features
self.model.classifier = nn.Linear(num_features, output_layer)
def forward(self, x):
"""
In the forward function we accept a Tensor of input data and we must return
a Tensor of output data. We can use Modules defined in the constructor as
well as arbitrary operators on Tensors.
"""
y_pred = self.model(x)
return y_pred
def get_model(model_name, output_layer = 5):
return mobilenet(model_name, output_layer=output_layer)