def main(img_path): img = cv2.imread(img_path) if img is None: print("invalid image!") exit(1) if False == exists('models'): print('no pretrained model found!') exit(1) data = tf.expand_dims(img.astype('float32') / 255., axis=0) # data.shape = (1, h, w, 3) if exists(join('models', 'deeplabv3plus_weights.h5')): deeplabv3plus = DeeplabV3Plus(1 + 80) deeplabv3plus.load_weights(join('models', 'deeplabv3plus_weights.h5')) else: deeplabv3plus = tf.keras.models.load_model(join( 'models', 'deeplabv3plus.h5'), custom_objects={'tf': tf}) preds = deeplabv3plus(data) # preds.shape = (1, h, w, 1 + 80) seg = tf.argmax(preds[0:1, ...], axis=-1) # cls.shape = (1, 256, 256) classes, _ = tf.unique(tf.reshape(seg, (-1, ))) # cls.shape = (class num) palette = tf.random.uniform(maxval=256, shape=(classes.shape[0], 3), dtype=tf.int32) # palette.shape = (class num, 3) colormap = tf.cast(tf.gather_nd(palette, tf.expand_dims(seg, axis=-1)), dtype=tf.float32) # colormap.shape = (1, 255, 255, 3) seg_img = tf.cast(colormap, dtype=tf.uint8) cv2.imshow('origin', img) cv2.imshow('segmentation', seg_img[0].numpy()) cv2.waitKey()
def main(): # distributed strategy strategy = tf.distribute.MirroredStrategy(); # load dataset trainset_filenames = [join('trainset', filename) for filename in listdir('trainset')]; testset_filenames = [join('testset', filename) for filename in listdir('testset')]; trainset = tf.data.TFRecordDataset(trainset_filenames).map(parse_function).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE); testset = tf.data.TFRecordDataset(testset_filenames).map(parse_function).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE); with strategy.scope(): deeplabv3plus = DeeplabV3Plus(3, 80 + 1); deeplabv3plus.compile(optimizer = 'adam', loss = 'sparse_categorical_crossentropy', metrics = ['sparse_categorical_accuracy']); deeplabv3plus.fit(trainset, epochs = 100, validation_data = testset); deeplabv3plus.save('deeplabv3plus.h5');
def main(): deeplabv3plus = DeeplabV3Plus(80 + 1) optimizer = tf.keras.optimizers.Adam( tf.keras.optimizers.schedules.ExponentialDecay(1e-3, decay_steps=60000, decay_rate=0.5)) checkpoint = tf.train.Checkpoint(model=deeplabv3plus, optimizer=optimizer) checkpoint.restore(tf.train.latest_checkpoint('checkpoints')) if False == exists('models'): mkdir('models') deeplabv3plus.save(join('models', 'deeplabv3plus.h5')) deeplabv3plus.save_weights(join('models', 'deeplabv3plus_weights.h5')) deeplabv3plus.get_layer('resnet50').save_weights( join('models', 'resnet50.h5'))
def main(): # distributed strategy strategy = tf.distribute.MirroredStrategy(); # variables for replica with strategy.scope(): deeplabv3plus = DeeplabV3Plus(3, 80 + 1); loss_object = tf.keras.losses.SparseCategoricalCrossentropy(reduction = tf.keras.losses.Reduction.NONE); def compute_loss(labels, predictions): per_example_loss = loss_object(labels, predictions); return tf.nn.compute_average_loss(per_example_loss, global_batch_size = batch_size * strategy.num_replicas_in_sync); optimizer = tf.keras.optimizers.Adam(tf.keras.optimizers.schedules.ExponentialDecay(1e-3, decay_steps = 60000, decay_rate = 0.5)); checkpoint = tf.train.Checkpoint(model = deeplabv3plus, optimizer = optimizer); train_loss = tf.keras.metrics.Mean(name = 'train loss', dtype = tf.float32); train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name = 'train accuracy'); test_loss = tf.keras.metrics.Mean(name = 'test loss', dtype = tf.float32); test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name = 'test accuracy'); # load dataset trainset_filenames = [join('trainset', filename) for filename in listdir('trainset')]; testset_filenames = [join('testset', filename) for filename in listdir('testset')]; trainset = tf.data.TFRecordDataset(trainset_filenames).repeat(-1).map(parse_function).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE); testset = tf.data.TFRecordDataset(testset_filenames).repeat(-1).map(parse_function).shuffle(batch_size).batch(batch_size).prefetch(tf.data.experimental.AUTOTUNE); dist_trainset = strategy.experimental_distribute_dataset(trainset); dist_testset = strategy.experimental_distribute_dataset(testset); dist_trainset_iter = iter(dist_trainset); dist_testset_iter = iter(dist_testset); # checkpoint if False == exists('checkpoints'): mkdir('checkpoints'); checkpoint.restore(tf.train.latest_checkpoint('checkpoints')); # log log = tf.summary.create_file_writer('checkpoints'); # train step (per GPU) def train_step(inputs): data, labels = inputs; with tf.GradientTape() as tape: predictions = deeplabv3plus(data, training = True); loss = compute_loss(labels, predictions); gradients = tape.gradient(loss, deeplabv3plus.trainable_variables); if tf.math.reduce_any([tf.math.reduce_any(tf.math.logical_or(tf.math.is_nan(grad), tf.math.is_inf(grad))) for grad in gradients]) == True: print('detected nan in grads, skip current iterations'); return 0; optimizer.apply_gradients(zip(gradients, deeplabv3plus.trainable_variables)); train_loss.update_state(loss); train_accuracy.update_state(labels, predictions); return loss; def test_step(inputs): data, labels = inputs; predictions = deeplabv3plus(data, training = False); loss = loss_object(labels, predictions); test_loss.update_state(loss); test_accuracy.update_state(labels, predictions); # train while True: samples = next(dist_trainset_iter); # every GPU runs its replica per_replica_losses = strategy.run(train_step, args = (samples,)); # reduce losses from all replica strategy.reduce(tf.distribute.ReduceOp.MEAN, per_replica_losses, axis = None); if tf.equal(optimizer.iterations % 100, 0): # save checkpoint checkpoint.save(join('checkpoints', 'ckpt')); # print test status for i in range(10): samples = next(dist_testset_iter); strategy.run(test_step, args = (samples,)); # write log with log.as_default(): tf.summary.scalar('train loss', train_loss.result(), step = optimizer.iterations); tf.summary.scalar('train accuracy', train_accuracy.result(), step = optimizer.iterations); tf.summary.scalar('test loss', test_loss.result(), step = optimizer.iterations); tf.summary.scalar('test accuracy', test_accuracy.result(), step = optimizer.iterations); print('Step #%d Train Loss: %.6f Train Accuracy: %.6f Test Loss: %.6f Test Accuracy: %.6f' % \ (optimizer.iterations, train_loss.result(), train_accuracy.result(), test_loss.result(), test_accuracy.result())); train_loss.reset_states(); train_accuracy.reset_states(); test_loss.reset_states(); test_accuracy.reset_states(); # break condition if train_loss.result() < 0.01: break; deeplabv3plus.save('deeplabv3plus.h5');
def main(): deeplabv3plus = DeeplabV3Plus(80 + 1) optimizer = tf.keras.optimizers.Adam( tf.keras.optimizers.schedules.ExponentialDecay(1e-3, decay_steps=110000, decay_rate=0.9)) checkpoint = tf.train.Checkpoint(model=deeplabv3plus, optimizer=optimizer) train_loss = tf.keras.metrics.Mean(name='train loss', dtype=tf.float32) train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy( name='train accuracy') test_loss = tf.keras.metrics.Mean(name='test loss', dtype=tf.float32) test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy( name='test accuracy') trainset_filenames = [ join('trainset', filename) for filename in listdir('trainset') ] testset_filenames = [ join('testset', filename) for filename in listdir('testset') ] trainset = tf.data.TFRecordDataset(trainset_filenames).repeat(-1).map( parse_function).shuffle(batch_size).batch(batch_size).prefetch( tf.data.experimental.AUTOTUNE) testset = tf.data.TFRecordDataset(testset_filenames).repeat(-1).map( parse_function).shuffle(batch_size).batch(batch_size).prefetch( tf.data.experimental.AUTOTUNE) trainset_iter = iter(trainset) testset_iter = iter(testset) # checkpoint if False == exists('checkpoints'): mkdir('checkpoints') checkpoint.restore(tf.train.latest_checkpoint('checkpoints')) # log log = tf.summary.create_file_writer('checkpoints') # train while True: data, labels = next(trainset_iter) err_labels = tf.boolean_mask( labels, tf.math.logical_or(tf.math.less(labels, 0), tf.math.greater(labels, 80))) if err_labels.shape[0] != 0: print('encounter invalid labels', err_labels) continue with tf.GradientTape() as tape: if tf.math.reduce_any( tf.math.logical_or(tf.math.is_nan(data), tf.math.is_inf(data))) == True: print('detected nan in data, skip current iterations') continue preds = deeplabv3plus(data, training=True) if tf.math.reduce_any( tf.math.logical_or(tf.math.is_nan(preds), tf.math.is_inf(preds))) == True: print('detected nan in preds, skip current iterations') pdb.set_trace() continue loss = tf.keras.losses.SparseCategoricalCrossentropy()(labels, preds) if tf.math.reduce_any( tf.math.logical_or(tf.math.is_nan(loss), tf.math.is_inf(loss))) == True: print('detected nan in loss, skip current iterations') continue grads = tape.gradient(loss, deeplabv3plus.trainable_variables) if tf.math.reduce_any([ tf.math.reduce_any( tf.math.logical_or(tf.math.is_nan(grad), tf.math.is_inf(grad))) for grad in grads ]) == True: print('detected nan in grads, skip current iterations') continue optimizer.apply_gradients(zip(grads, deeplabv3plus.trainable_variables)) train_loss.update_state(loss) train_accuracy.update_state(labels, preds) if tf.equal(optimizer.iterations % 10000, 0): # save checkpoint checkpoint.save(join('checkpoints', 'ckpt')) if tf.equal(optimizer.iterations % 1000, 0): # evaluate for i in range(10): data, labels = next(testset_iter) err_labels = tf.boolean_mask( labels, tf.math.logical_or(tf.math.less(labels, 0), tf.math.greater(labels, 80))) if err_labels.shape[0] != 0: print('encounter invalid labels', err_labels) continue if tf.math.reduce_any( tf.math.logical_or(tf.math.is_nan(data), tf.math.is_inf(data))) == True: print('detected nan in data, skip current iterations') continue preds = deeplabv3plus(data, training=False) loss = tf.keras.losses.SparseCategoricalCrossentropy()(labels, preds) test_loss.update_state(loss) test_accuracy.update_state(labels, preds) # write log with log.as_default(): tf.summary.scalar('train loss', train_loss.result(), step=optimizer.iterations) tf.summary.scalar('train accuracy', train_accuracy.result(), step=optimizer.iterations) tf.summary.scalar('test loss', test_loss.result(), step=optimizer.iterations) tf.summary.scalar('test accuracy', test_accuracy.result(), step=optimizer.iterations) seg = tf.argmax(preds[0:1, ...], axis=-1) # cls.shape = (1, 256, 256) classes, _ = tf.unique(tf.reshape(seg, (-1, ))) # cls.shape = (class num) palette = tf.random.uniform(maxval=256, shape=(classes.shape[0], 3), dtype=tf.int32) # palette.shape = (class num, 3) colormap = tf.cast(tf.gather_nd(palette, tf.expand_dims(seg, axis=-1)), dtype=tf.float32) # colormap.shape = (1, 255, 255, 3) img = tf.cast(tf.clip_by_value( tf.math.rint(0.5 * colormap + 0.5 * data[0:1, ..., ::-1] * 255.), 0, 255), dtype=tf.uint8) tf.summary.image('segmentation', img, step=optimizer.iterations) print('Step #%d Train Loss: %.6f Train Accuracy: %.6f Test Loss: %.6f Test Accuracy: %.6f' % \ (optimizer.iterations, train_loss.result(), train_accuracy.result(), test_loss.result(), test_accuracy.result())) # break condition if train_loss.result() < 0.01: break # reset train_loss.reset_states() train_accuracy.reset_states() test_loss.reset_states() test_accuracy.reset_states() deeplabv3plus.save('deeplabv3plus.h5')