示例#1
0
def train(create_model, X, y, batch_size, epochs, gpu_count, parameter_server,
          method):
    if gpu_count > 1:
        ps_device = '/gpu:0' if parameter_server == 'gpu' else '/cpu:0'

        with tf.device(ps_device):
            serial_model = create_model()

        if method == 'kuza55':
            from keras_tf_multigpu.kuza55 import make_parallel
            model = make_parallel(serial_model,
                                  gpu_count=gpu_count,
                                  ps_device=ps_device)
        elif method == 'avolkov1':
            from keras_tf_multigpu.avolkov1 import make_parallel, get_available_gpus
            gpus_list = get_available_gpus(gpu_count)
            model = make_parallel(serial_model,
                                  gdev_list=gpus_list,
                                  ps_device=ps_device)
        elif method == 'fchollet':
            # requires Keras (2.0.9?) https://github.com/fchollet/keras/commit/3dd3e8331677e68e7dec6ed4a1cbf16b7ef19f7f
            from keras.utils import multi_gpu_model
            model = multi_gpu_model(serial_model, gpus=gpu_count)
    else:
        model = serial_model = create_model()

    print('Number of parameters:', serial_model.count_params())

    model.compile(loss='categorical_crossentropy',
                  optimizer='sgd',
                  metrics=['accuracy'])
    gauge = SamplesPerSec(batch_size)
    model.fit(X, y, batch_size=batch_size, epochs=epochs, callbacks=[gauge])
    gauge.print_results()
    model.add(Flatten())
    model.add(Dense(512))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(num_classes))
    model.add(Activation('softmax'))

    return model

with tf.device(ps_device):
    serial_model = basic_model()
    print('Serial model:')
    serial_model.summary()

    model = make_parallel(tower, gpu_count, ps_device)
    print('Multi-GPU model:')
    model.summary()

# initiate RMSprop optimizer
opt = keras.optimizers.rmsprop(lr=0.0001, decay=1e-6)

# Let's train the model using RMSprop
model.compile(loss='categorical_crossentropy',
              optimizer=opt,
              metrics=['accuracy'])

x_train = x_train.astype('float32')
x_test = x_test.astype('float32')
x_train /= 255
x_test /= 255
def cnn(channel=3):
    """
    Architecture and model of the CNN.

    :return: object, model of the CNN.
    """
    sgd = optimizers.SGD(lr=0.01, momentum=0.0, decay=0.0, nesterov=True)
    inputs = Input(shape=(40, 80, 1))
    gaussian = TruncatedNormal(stddev=0.01, seed=None)

    pitch = Conv2D(filters=32,
                   kernel_size=(32, 1),
                   kernel_initializer=gaussian,
                   activation='relu',
                   name='conv_1')(inputs)
    pitch = BatchNormalization()(pitch)
    pitch = MaxPooling2D(pool_size=(1, 80))(pitch)
    pitch = Reshape((1, 9, -1))(pitch)

    tempo = Conv2D(filters=32,
                   kernel_size=(1, 60),
                   kernel_initializer=gaussian,
                   activation='relu',
                   name='conv_2')(inputs)
    tempo = BatchNormalization()(tempo)
    tempo = MaxPooling2D(pool_size=(40, 1))(tempo)

    bass = Conv2D(filters=32,
                  kernel_size=(13, 9),
                  kernel_initializer=gaussian,
                  activation='relu',
                  name='conv_3')(inputs)
    bass = BatchNormalization()(bass)
    bass = MaxPooling2D(pool_size=(4, 4))(bass)
    bass = Reshape((1, 126, -1))(bass)

    if channel == 2:
        concatenate = Concatenate(axis=2)([pitch, tempo])
        flatten = Flatten()(concatenate)
        dense = Dense(200,
                      kernel_initializer=gaussian,
                      activation='relu',
                      name='dense_1')(flatten)
    else:  # channel == 3
        concatenate = Concatenate(axis=2)([pitch, tempo, bass])
        flatten = Flatten()(concatenate)
        dense = Dense(400,
                      kernel_initializer=gaussian,
                      activation='relu',
                      name='dense_1')(flatten)

    dropout = Dropout(0.5)(dense)
    predictions = Dense(10,
                        kernel_initializer=gaussian,
                        activation='softmax',
                        name='dense_2')(dropout)

    model = Model(inputs=inputs, outputs=predictions)
    model = make_parallel(model, gpu_count=2)
    model.compile(optimizer=sgd,
                  loss='categorical_crossentropy',
                  metrics=['accuracy'])

    plot_model(model,
               to_file='./plots/model_plot.png',
               show_shapes=True,
               show_layer_names=True)

    return model
    model = Model(inputs=input, outputs=output)

    print('Single tower model:')
    model.summary()
    return model

tensorboard_dir = './tensorboard-logs/mnist_mlp_multi_ps_cpu/%d-gpu_%s' \
    % (gpu_count, os.environ.get('CUDA_VISIBLE_DEVICES', ''))

with tf.device(ps_device):
    serial_model = basic_model()
    print('Serial model:')
    serial_model.summary()

    model = make_parallel(serial_model, gpu_count, ps_device)
    print('Multi-GPU model:')
    model.summary()

    model.compile(loss='categorical_crossentropy',
                  optimizer=RMSprop(),
                  metrics=['accuracy'])

    summary_writer = tf.summary.FileWriter(tensorboard_dir, sess.graph)
    summary_writer.flush()

    tensorboard_cb = keras.callbacks.TensorBoard(log_dir=tensorboard_dir)
    history = model.fit(
        x_train,
        y_train,
        batch_size=batch_size,