AIProjekt/net_training.py
Jakub Adamski f6c4f94b0c training
2020-05-18 10:55:02 +02:00

36 lines
1.3 KiB
Python

import matplotlib.pyplot as plt
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet import preprocess_input, decode_predictions
from tensorflow.keras.layers import Dense
from tensorflow.keras.models import Model
import numpy as np
import os, random
#podstawa modelu
base_model = MobileNetV2(include_top=False, weights="imagenet", pooling='avg')
#model
x=base_model.output
preds=Dense(4,activation='softmax')(x)
model=Model(inputs=base_model.input,outputs=preds)
#tylko ostatnie 20 warstw uczymy
for layer in model.layers[:20]:
layer.trainable=False
for layer in model.layers[20:]:
layer.trainable=True
#generator obrazkow
train_datagen=ImageDataGenerator(preprocessing_function=preprocess_input)
train_generator=train_datagen.flow_from_directory('./dataset', target_size=(224,224), color_mode='rgb', batch_size=32, class_mode='categorical', shuffle=True)
#uczenie //to dzielenie i podloga
model.compile(optimizer='Adam',loss='categorical_crossentropy',metrics=['accuracy'])
step_size_train=train_generator.n//train_generator.batch_size
model.fit_generator(generator=train_generator, steps_per_epoch=step_size_train, epochs=10)
#zapis
model.save('moj_model.h5')