Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- from tensorflow.keras.layers import Conv2D, AveragePooling2D, GlobalAveragePooling2D
- from tensorflow.keras.models import Sequential
- from tensorflow.keras.layers import Dense
- from argparse import ArgumentParser
- from tensorflow.keras import utils
- from datetime import datetime
- import tensorflow_hub as hub
- from PIL import ImageFile
- import tensorflow as tf
- import pandas as pd
- import numpy as np
- import itertools
- import PIL
- import sys
- import pdb
- import os
- ImageFile.LOAD_TRUNCATED_IMAGES = True
- parser = ArgumentParser()
- parser.add_argument('--csv_file', '-c',
- help='path to csv file')
- parser.add_argument('--gpus', '-g', nargs='+', type=str, default=["0"],
- help='number of GPUs to train on')
- parser.add_argument('--training_type', '-tt', choices=['gender', 'age', 'mask', 'age_gender'],
- help='type of classifier to train')
- parser.add_argument('--count', type=int, default=5000,
- help='Number of images to be selected from each age group')
- args = parser.parse_args()
- def build_model_from_keras(train_generator):
- base_model = tf.keras.applications.MobileNetV2(include_top=False, weights=None, input_shape=(96,96,3), pooling='avg')
- optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)
- if args.training_type != 'age_gender':
- prediction = tf.keras.layers.Dense(units=len(train_generator.class_indices),
- activation='softmax', name='pred')(base_model.output)
- model = tf.keras.models.Model(inputs=base_model.input, outputs=prediction)
- model.compile(
- optimizer=optimizer,
- loss=tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1),
- metrics=['accuracy'])
- else:
- dense_age = tf.keras.layers.Dense(units=64, activation=tf.keras.layers.ReLU(6.0), name='dense_age')(base_model.output)
- dense_gender = tf.keras.layers.Dense(units=64, activation=tf.keras.layers.ReLU(6.0), name='dense_gender')(base_model.output)
- dropout_age = tf.keras.layers.Dropout(rate=0.2)(dense_age)
- dropout_gender = tf.keras.layers.Dropout(rate=0.2)(dense_gender)
- pred_age = tf.keras.layers.Dense(units=3, activation='softmax', name='pred_age')(dropout_age)
- pred_gender = tf.keras.layers.Dense(units=2, activation='softmax', name='pred_gender')(dropout_gender)
- model = tf.keras.models.Model(inputs=base_model.input, outputs=[pred_age, pred_gender])
- model.compile(
- optimizer=optimizer,
- loss=[tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1), tf.keras.losses.CategoricalCrossentropy(from_logits=True, label_smoothing=0.1)],
- metrics=[tf.keras.metrics.Accuracy(name='age_accuracy'), tf.keras.metrics.Accuracy(name='gender_accuracy')])
- return model
- ###############################################################################
- ### Data ##############################################################
- ###############################################################################
- df = pd.read_csv(args.csv_file)
- if args.training_type in ['age', 'age_gender']:
- df = df[df['mask'] == 'f']
- df_18_25 = df.loc[df['age_group'] == '(18, 25)']
- df_26_39 = df.loc[df['age_group'] == '(26, 39)']
- df_40_55 = df.loc[df['age_group'] =='(40, 55)']
- df_18_25 = df_18_25.sample(args.count, random_state=42, replace=True)
- df_26_39 = df_26_39.sample(args.count, random_state=42, replace=True)
- df_40_55 = df_40_55.sample(args.count, random_state=42, replace=True)
- df = pd.concat([df_18_25,df_26_39,df_40_55])
- y_col = 'age_group'
- if args.training_type == 'age_gender':
- y_col = ['age_group', 'gender']
- df = df.sample(frac=1, random_state=42)
- val_idx = round(len(df) * 0.2)
- val_data = df.iloc[-val_idx:]
- train_data = df.iloc[:-val_idx]
- print(df)
- datagen_kwargs = dict(
- rescale=1./255,
- fill_mode='constant',
- cval=127)
- dataflow_kwargs = dict(
- target_size=(96,96),
- batch_size=32,
- interpolation="bilinear")
- valid_datagen = tf.keras.preprocessing.image.ImageDataGenerator(**datagen_kwargs)
- valid_generator = valid_datagen.flow_from_dataframe(
- directory=None,
- dataframe=val_data,
- x_col='filename',
- y_col=y_col,
- class_mode='categorical' if args.training_type != 'age_gender' else 'multi_output',
- shuffle=False,
- **dataflow_kwargs)
- train_datagen = tf.keras.preprocessing.image.ImageDataGenerator(
- rotation_range=45,
- horizontal_flip=True,
- width_shift_range=0.2,
- height_shift_range=0.2,
- shear_range=40,
- zoom_range=(0.5,1.5),
- brightness_range=(0.1,0.9),
- channel_shift_range=75,
- **datagen_kwargs)
- train_generator = train_datagen.flow_from_dataframe(
- directory=None,
- dataframe=train_data,
- x_col='filename',
- y_col=y_col,
- class_mode='categorical' if args.training_type != 'age_gender' else 'multi_output',
- shuffle=False,
- **dataflow_kwargs)
- try:
- print(f"class indices: {train_generator.class_indices}")
- except Exception as e:
- print(f"could not print class indices")
- try:
- print(f"num classes: {train_generator.num_classes}")
- except Exception as e:
- print(f"could not print num classes")
- ###############################################################################
- ### Model Definition ##################################################
- ###############################################################################
- physical_devices = tf.config.list_physical_devices('GPU')
- tf.config.set_visible_devices(physical_devices[int(args.gpus[0])], 'GPU')
- model=build_model_from_keras(train_generator)
- for l in model.layers:
- l.trainable = True
- model.trainable=True
- model.summary()
- print(f"Train Data:\n{train_data.shape}\n")
- print(f"Validation Data:\n{val_data.shape}\n")
- train_date = None
- ###############################################################################
- ### Training ##########################################################
- ###############################################################################
- train_date = datetime.now().strftime("%Y%m%d-%H%M%S")
- save_path = os.path.join('mobilenetv2', train_date)
- os.makedirs(save_path)
- model_saver = tf.keras.callbacks.ModelCheckpoint(os.path.join(
- save_path, '{epoch:02d}-{accuracy:.3f}-{val_accuracy:.3f}.h5'), save_freq='epoch', verbose=1)
- callbacks = [model_saver]
- steps_per_epoch = train_generator.samples // train_generator.batch_size
- validation_steps = valid_generator.samples // valid_generator.batch_size
- history = model.fit_generator(
- generator=train_generator,
- callbacks=callbacks,
- epochs=1,
- steps_per_epoch=steps_per_epoch,
- validation_data=valid_generator,
- validation_steps=validation_steps)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement