def model_choose(ih, iw, nb_conv, size_conv, nb_down, nb_gpu): if nb_down == 3: mdl = transformer3_pooling(ih, iw, nb_conv, size_conv, nb_gpu) elif nb_down == 1: mdl = transformer_3CNN(ih, iw, nb_conv, size_conv, nb_gpu) else: mdl = transformer2(ih, iw, nb_conv, size_conv, nb_gpu) return mdl
def model_choose(ih, iw, conv_nb, conv_size, model_layers, gpu_nb): if model_layers == 3: mdl = transformer3_pooling(ih, iw, conv_nb, conv_size, gpu_nb) else: mdl = transformer2(ih, iw, conv_nb, conv_size, gpu_nb) return mdl