def model_def()

in src/train.py [0:0]

• 24 lines of code
• 5 McCabe index (conditional complexity)

def model_def(learning_rate, mpi=False, hvd=False):
    inputs = Input(shape=(HEIGHT, WIDTH, DEPTH), name='inputs')
    base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(HEIGHT, WIDTH, DEPTH),
                          input_tensor=inputs)
    base_model.trainable = False

    x1 = base_model.output
    x1 = GlobalAveragePooling2D()(x1)
    x1 = Dropout(0.5)(x1)
    x1 = Dense(512, activation='relu')(x1)
    x1 = Dropout(0.5)(x1)
    x1 = Dense(NUM_CLASSES, activation='softmax')(x1)

    model = Model(inputs=[base_model.input], outputs=[x1])
    for layer in model.layers[:20]:
        layer.trainable = False
    for layer in model.layers[20:]:
        layer.trainable = True

    size = hvd.size() if mpi else 1

    opt = SGD(lr=learning_rate * size)

    if mpi:
        opt = hvd.DistributedOptimizer(opt)
    model.compile(loss='categorical_crossentropy',
                  optimizer=opt,
                  metrics=['accuracy'], experimental_run_tf_function=False)
    return model