forked from XuChunqiao/Tensorflow-Fast-MPNCOV
-
Notifications
You must be signed in to change notification settings - Fork 0
/
model_init.py
57 lines (48 loc) · 2.09 KB
/
model_init.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
from src.network import *
import tensorflow as tf
import numpy as np
import warnings
__all__ = ['Newmodel', 'get_model']
class Newmodel(Basemodel):
"""replace the image representation method and classifier
Args:
modeltype: model archtecture
representation: image representation method
num_classes: the number of classes
freezed_layer: the end of freezed layers in network
pretrained: whether use pretrained weights or not
"""
def __init__(self, modeltype, representation, num_classes, freezed_layer, pretrained=False):
super(Newmodel, self).__init__(modeltype, pretrained)
if representation is not None:
representation_method = representation['function']
representation.pop('function')
representation_args = representation
representation_args['input_dim'] = self.representation_dim
self.representation = representation_method(**representation_args)
if not self.pretrained:
if modeltype.startswith('vgg'):
self.classifier.pop()
self.classifier.add(tf.keras.layers.Dense(num_classes))
else:
self.classifier = tf.keras.layers.Dense(num_classes)
else:
self.classifier = tf.keras.layers.Dense(num_classes)
else:
if modeltype.startswith('vgg'):
self.classifier.pop()
self.classifier.add(tf.keras.layers.Dense(num_classes))
else:
self.classifier = tf.keras.layers.Dense(num_classes)
if freezed_layer:
self.features.trainable = False
self.representation.trainable = False
def get_model(modeltype, representation, num_classes, freezed_layer, pretrained=False):
_model = Newmodel(modeltype, representation, num_classes, freezed_layer, pretrained=pretrained)
input = tf.random.normal([1,224,224,3])
_model(input, training=False)
_model.features.summary()
# _model.representation
# _model.classifier.summary()
_model.summary()
return _model