TensorFlow培训每批速度变慢

时间:2020-11-08 11:09:25

标签: performance optimization deep-learning tensorflow2.0 loss-function

我是TensorFlow的新手,并且可以通过修改官方网站上的教程来成功运行代码。

我在StackOverflow上检查了其他一些答案,这表示我的问题很可能是由于每次都在图形中添加了一些内容。但是,我不知道在哪里可以找到可能引起这种情况的代码。

我还使用tf.py_function映射数据集,因为我确实需要在映射中启用热切模式。

def get_dataset(data_index):
    # data_index is a Pandas Dataframe that contains image/label pair info, each row is one pair
    data_index = prepare_data_index(data_index)
    # shuffle dataframe here because dataset.shuffle is taking very long time.
    data_index = data_index.sample(data_index.shape[0])

    path = path_to_img_dir

    # list of dataframe indices indicating rows that are going to be included in the dataset for training. 
    indices_ls = ['{}_L'.format(x) for x in list(data_index.index)] + ['{}_R'.format(x) for x in list(data_index.index)]
  
    # around 310k images
    image_count = len(indices_ls)

    list_ds = tf.data.Dataset.from_tensor_slices(indices_ls)
    # dataset.shuffle is commented out because it takes too much time
    # list_ds = list_ds.shuffle(image_count, reshuffle_each_iteration=False)

    val_size = int(image_count * 0.2)
    train_ds = list_ds.skip(val_size)
    val_ds = list_ds.take(val_size)

    def get_label(index):
        index = str(np.array(index).astype(str))
        delim = index.split('_')
        state = delim[1]
        index = int(delim[0])

        if state == 'R':
            label = data_index.loc[index][right_labels].to_numpy().flatten()
        elif state == 'L':
            label = data_index.loc[index][left_labels].to_numpy().flatten()

        return tf.convert_to_tensor(label , dtype=tf.float16)

    def get_img(index):
        index = str(np.array(index).astype(str))
        delim = index.split('_')
        state = delim[1]
        index = int(delim[0])

        file_path = '{}_{}.jpg'.format(data_index.loc[index, 'sub_folder'],
                                   str(int(data_index.loc[index, 'img_index'])).zfill(4)
                                   )
        img = tf.io.read_file(os.path.join(path, file_path))
        img = tf.image.decode_jpeg(img, channels=3)
        full_width = 320
        img = tf.image.resize(img, [height, full_width])

        # Crop half of the image depending on the state
        if state == 'R':
            
            img = tf.image.crop_to_bounding_box(img, offset_height=0, offset_width=0, target_height=height,
                                            target_width=int(full_width / 2))
            img = tf.image.flip_left_right(img)
        elif state == 'L':
            
            img = tf.image.crop_to_bounding_box(img, offset_height=0, offset_width=int(full_width / 2), target_height=height,
                                            target_width=int(full_width / 2))
        img = tf.image.resize(img, [height, width])
        img = tf.keras.preprocessing.image.array_to_img(
        img.numpy(), data_format=None, scale=True, dtype=None
    )
        # Apply auto white balancing, output an np array
        img = AWB(img)
        img = tf.convert_to_tensor(img, dtype=tf.float16)

        return img

    def process_path(index):
        label = get_label(index)
        img = get_img(index)
        return img, label

    AUTOTUNE = tf.data.experimental.AUTOTUNE

    train_ds = train_ds.map(lambda x: tf.py_function(
    process_path,
    [x], (tf.float16, tf.float16)), num_parallel_calls=AUTOTUNE)
    val_ds = val_ds.map(lambda x: tf.py_function(
    process_path,
    [x], (tf.float16, tf.float16)), num_parallel_calls=AUTOTUNE)

    def configure_for_performance(ds):
        ds = ds.cache()
        # ds = ds.shuffle(buffer_size=image_count)
        ds = ds.batch(batch_size)
        ds = ds.prefetch(buffer_size=AUTOTUNE)
        return ds

    train_ds = configure_for_performance(train_ds)
    val_ds = configure_for_performance(val_ds)

    return train_ds, val_ds

有人可以帮助我吗?谢谢!

这是我其余的代码。

def initialize_model():
    IMG_SIZE = (height, width)

    preprocess_input = tf.keras.applications.vgg19.preprocess_input

    IMG_SHAPE = IMG_SIZE + (3,)
    base_model = tf.keras.applications.VGG19(input_shape=IMG_SHAPE,
                                         include_top=False,
                                         weights='imagenet')

    global_average_layer = tf.keras.layers.GlobalAveragePooling2D()
    prediction_layer = tf.keras.layers.Dense(class_num, activation=tf.nn.sigmoid, use_bias=True)

    inputs = tf.keras.Input(shape=(height, width, 3))
    x = preprocess_input(inputs)
    x = base_model(x, training=True)
    x = global_average_layer(x)
    outputs = prediction_layer(x)
    model = tf.keras.Model(inputs, outputs)

    def custom_loss(y_gt, y_pred):
        
        L1_loss_out = tf.math.abs(tf.math.subtract(y_gt, y_pred))
        scaler = tf.pow(50.0, y_gt)
        scaled_loss = tf.math.multiply(L1_loss_out, scaler)
        scaled_loss = tf.math.reduce_mean(
            scaled_loss, axis=None, keepdims=False, name=None
        )
        return scaled_loss

    base_learning_rate = 0.001
    model.compile(optimizer=tf.keras.optimizers.SGD(learning_rate=base_learning_rate, momentum=0.9),
                  loss=custom_loss,
                  metrics=['mean_absolute_error']
                  )

    return model


def train(data_index, epoch_num, save_path):
    train_dataset, validation_dataset = get_dataset(data_index)
    model = initialize_model()
    model.summary()

    history = model.fit(train_dataset,
                        epochs=epoch_num,
                        validation_data=validation_dataset)
    model.save_weights(save_path)

    return model, history

0 个答案:

没有答案