File size: 1,429 Bytes
65e71e9 |
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 |
from tensorflow import keras
class FCNutr(keras.Model):
def __init__(self, nutr_names, crop_size, hidden_dim, num_shared_hidden, use_dropout):
super().__init__()
self.cnn = keras.applications.InceptionV3(
include_top=False,
weights="imagenet",
input_shape=crop_size + (3,),
pooling='avg'
)
shared_layers = []
for i in range(num_shared_hidden):
if use_dropout:
shared_layers.append(keras.layers.Dropout(0.2 if i == 0 else 0.5))
shared_layers.append(keras.layers.Dense(hidden_dim, activation="relu"))
self.shared = keras.Sequential(shared_layers, name='shared')
self.multitask_heads = []
for name in nutr_names:
head_layers = [
keras.layers.Dense(hidden_dim, activation="relu"),
keras.layers.Dense(1, activation="relu")
]
if use_dropout:
head_layers.insert(0, keras.layers.Dropout(0.5))
head_layers.insert(2, keras.layers.Dropout(0.5))
self.multitask_heads.append(keras.Sequential(head_layers, name=name))
def call(self, inputs, training=False):
x = self.cnn(inputs, training=training)
x = self.shared(x, training=training)
return {head.name: head(x, training=training) for head in self.multitask_heads}
|