80 lines
2.3 KiB
Python
80 lines
2.3 KiB
Python
#!/usr/bin/env python
|
|
# coding: utf-8
|
|
#
|
|
# wut-worker.py --- Runs on worker nodes.
|
|
#
|
|
# Start with wut-worker shell script to set correct
|
|
# environmental variables.
|
|
|
|
from __future__ import absolute_import, division, print_function, unicode_literals
|
|
import simplejson as json
|
|
import os
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
import tensorflow.python.keras
|
|
from tensorflow.python.keras import Sequential
|
|
from tensorflow.python.keras.layers import Activation, Dropout, Flatten, Dense
|
|
from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
|
|
from tensorflow.python.keras.layers import Convolution2D, MaxPooling2D, ZeroPadding2D
|
|
from tensorflow.python.keras import optimizers
|
|
from tensorflow.python.keras.preprocessing import image
|
|
from tensorflow.python.keras.models import load_model
|
|
from tensorflow.python.keras.preprocessing.image import load_img
|
|
from tensorflow.python.keras.preprocessing.image import img_to_array
|
|
from tensorflow.keras.layers import Dense, Conv2D, Flatten, Dropout, MaxPooling2D
|
|
from tensorflow.python.keras.models import Model
|
|
from tensorflow.python.keras.layers import Input, concatenate
|
|
|
|
#batch_size = 32
|
|
#epochs = 4
|
|
IMG_HEIGHT = 416
|
|
IMG_WIDTH= 804
|
|
|
|
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()
|
|
|
|
|
|
def get_uncompiled_model():
|
|
model = Sequential([
|
|
Conv2D(16, 3, padding='same', activation='relu', input_shape=(IMG_HEIGHT, IMG_WIDTH ,3)),
|
|
MaxPooling2D(),
|
|
Conv2D(32, 3, padding='same', activation='relu'),
|
|
MaxPooling2D(),
|
|
Conv2D(64, 3, padding='same', activation='relu'),
|
|
MaxPooling2D(),
|
|
Flatten(),
|
|
Dense(512, activation='relu'),
|
|
Dense(1, activation='sigmoid')
|
|
])
|
|
return model
|
|
|
|
def get_compiled_model():
|
|
model = get_uncompiled_model()
|
|
model.compile(optimizer='adam',
|
|
loss='binary_crossentropy',
|
|
metrics=['accuracy'])
|
|
return model
|
|
|
|
def get_fit_model():
|
|
model = get_compiled_model()
|
|
model.fit(
|
|
model )
|
|
return model
|
|
|
|
#def get_fit_model():
|
|
# model = get_compiled_model()
|
|
# model.fit(
|
|
# train_data_gen,
|
|
# steps_per_epoch=total_train // batch_size,
|
|
# epochs=epochs,
|
|
# validation_data=val_data_gen,
|
|
# validation_steps=total_val // batch_size,
|
|
# verbose=2
|
|
# )
|
|
# return model
|
|
|
|
with strategy.scope():
|
|
get_uncompiled_model()
|
|
get_compiled_model()
|
|
get_fit_model()
|
|
|