import pytorch_metric_learning.samplers

from models.conv_net import ConvNet
from models.deepeyedentification import DeepEyedentificationLive
from models.eye_know_you_too import EKY2


models = {
    'del_fast': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'del-fast-trafo',
        'preprocessing_kwargs': {
            'v_min': 40.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [32, 32, 32, 512, 512, 512, 512, 512, 512],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del_fast_vmin1': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'del-fast-trafo',
        'preprocessing_kwargs': {
            'v_min': 1.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [32, 32, 32, 512, 512, 512, 512, 512, 512],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del_fast_vmin4': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'del-fast-trafo',
        'preprocessing_kwargs': {
            'v_min': 4.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [32, 32, 32, 512, 512, 512, 512, 512, 512],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del_fast_vmin10': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'del-fast-trafo',
        'preprocessing_kwargs': {
            'v_min': 10.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [32, 32, 32, 512, 512, 512, 512, 512, 512],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del_fast_zstd': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [32, 32, 32, 512, 512, 512, 512, 512, 512],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del_slow': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'del-slow-trafo',
        'preprocessing_kwargs': {
            'c': 0.02,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del_slow_zstd': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_padding': 'same',
            'pool_before_bn_in_first_layer': True,
        },
    },
    'del': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged',
        'preprocessing_kwargs': {
            'c': 0.02,
            'v_min': 40.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast',
            'path_weights_fast_subnet': 'del_fast',
            'config_slow_subnet': 'del_slow',
            'path_weights_slow_subnet': 'del_slow',
        },
    },
    'del_vmin1': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged',
        'preprocessing_kwargs': {
            'c': 0.02,
            'v_min': 1.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast_vmin1',
            'path_weights_fast_subnet': 'del_fast_vmin1',
            'config_slow_subnet': 'del_slow',
            'path_weights_slow_subnet': 'del_slow',
        },
    },
    'del_vmin4': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged',
        'preprocessing_kwargs': {
            'c': 0.02,
            'v_min': 4.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast_vmin4',
            'path_weights_fast_subnet': 'del_fast_vmin4',
            'config_slow_subnet': 'del_slow',
            'path_weights_slow_subnet': 'del_slow',
        },
    },
    'del_vmin10': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged',
        'preprocessing_kwargs': {
            'c': 0.02,
            'v_min': 10.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast_vmin10',
            'path_weights_fast_subnet': 'del_fast_vmin10',
            'config_slow_subnet': 'del_slow',
            'path_weights_slow_subnet': 'del_slow',
        },
    },
    'del_single_stage': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged',
        'preprocessing_kwargs': {
            'c': 0.02,
            'v_min': 40.0,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast',
            'path_weights_fast_subnet': None,
            'config_slow_subnet': 'del_slow',
            'path_weights_slow_subnet': None,
        },
    },
    'del_zstd': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged-zstd',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast_zstd',
            'path_weights_fast_subnet': 'del_fast_zstd',
            'config_slow_subnet': 'del_slow_zstd',
            'path_weights_slow_subnet': 'del_slow_zstd',
        },
    },
    'del_zstd_tanh': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged-zstd-tanh',
        'preprocessing_kwargs': {
            'c': 0.02,
        },
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast_zstd',
            'path_weights_fast_subnet': 'del_fast_zstd',
            'config_slow_subnet': 'del_slow',
            'path_weights_slow_subnet': 'del_slow',
        },
    },
    'del_zstd_single_stage': {
        'class': DeepEyedentificationLive,
        'framework': 'pytorch',
        'preprocessing': 'del-merged-zstd',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'n_channels': '$n_channels',
            'sequence_length': '$sequence_length',
            'dense_units': [256, 128],
            'config_fast_subnet': 'del_fast_zstd',
            'path_weights_fast_subnet': None,
            'config_slow_subnet': 'del_slow_zstd',
            'path_weights_slow_subnet': None,
        },
    },
    'eky2': {
        'class': EKY2,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 256,
        'early_stopping': False,
        'sampler': pytorch_metric_learning.samplers.MPerClassSampler,
        'sampler_kwargs': {
            'm': 16,
            'batch_size': 256,
            'length_before_new_iter': 100000,
        },
        'init_kwargs': {
            'input_channels': '$n_channels',
            'output_units': '$n_classes',
            'add_dense_connections': True,
            'loss_weights_categorical': 0.1,
            'loss_weights_metric': 1.0,
            'default_forward_mode': 'class',
        },
    },
    'eky2_bs224_m32': {
        'class': EKY2,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 224,
        'early_stopping': False,
        'sampler': pytorch_metric_learning.samplers.MPerClassSampler,
        'sampler_kwargs': {
            'm': 32,
            'batch_size': 224,
            'length_before_new_iter': 100000,
        },
        'init_kwargs': {
            'input_channels': '$n_channels',
            'output_units': '$n_classes',
            'add_dense_connections': True,
            'loss_weights_categorical': 0.1,
            'loss_weights_metric': 1.0,
            'default_forward_mode': 'class',
        },
    },
    'eky2_noskip': {
        'class': EKY2,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 256,
        'early_stopping': False,
        'init_kwargs': {
            'input_channels': '$n_channels',
            'output_units': '$n_classes',
            'add_dense_connections': False,
            'loss_weights_categorical': 0.1,
            'loss_weights_metric': 1.0,
            'default_forward_mode': 'class',
        },
    },
    'eky2_multiclass': {
        'class': EKY2,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 256,
        'early_stopping': False,
        'sampler': pytorch_metric_learning.samplers.MPerClassSampler,
        'sampler_kwargs': {
            'm': 16,
            'batch_size': 256,
            'length_before_new_iter': 100000,
        },
        'init_kwargs': {
            'input_channels': '$n_channels',
            'output_units': '$n_classes',
            'add_dense_connections': True,
            'loss_weights_categorical': 1.0,
            'loss_weights_metric': 0.0,
            'default_forward_mode': 'class',
        },
    },
    'eky2_multiclass_noskip': {
        'class': EKY2,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 256,
        'early_stopping': False,
        'init_kwargs': {
            'input_channels': '$n_channels',
            'output_units': '$n_classes',
            'add_dense_connections': False,
            'loss_weights_categorical': 1.0,
            'loss_weights_metric': 0.0,
            'default_forward_mode': 'class',
        },
    },

    'cnn1_dense1_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128],
            'conv_filters': [32],
            'conv_sizes': [5],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1],
            'pool_sizes': [2],
            'pool_strides': [2],

        },
    },
    'cnn2_dense1_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128],
            'conv_filters': [32, 32],
            'conv_sizes': [5, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1],
            'pool_sizes': [2, 2],
            'pool_strides': [2, 2],

        },
    },
    'cnn3_dense2_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128, 128],
            'conv_filters': [32, 32, 32],
            'conv_sizes': [5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1],
            'pool_sizes': [2, 2, 2],
            'pool_strides': [2, 2, 2],

        },
    },
    'cnn6_dense2_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128, 128],
            'conv_filters': [32, 32, 32, 32, 32, 32],
            'conv_sizes': [5, 5, 5, 3, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2],

        },
    },
    'cnn6_dense3_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2],

        },
    },
    'cnn7_dense2_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 128],
            'conv_filters': [128, 128, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn8_dense2_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn9_dense2_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn9_dense3_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn10_dense3_notrafo_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': False,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn1_dense1_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128],
            'conv_filters': [32],
            'conv_sizes': [5],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1],
            'pool_sizes': [2],
            'pool_strides': [2],

        },
    },
    'cnn2_dense1_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128],
            'conv_filters': [32, 32],
            'conv_sizes': [5, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1],
            'pool_sizes': [2, 2],
            'pool_strides': [2, 2],

        },
    },
    'cnn3_dense2_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128, 128],
            'conv_filters': [32, 32, 32],
            'conv_sizes': [5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1],
            'pool_sizes': [2, 2, 2],
            'pool_strides': [2, 2, 2],

        },
    },
    'cnn6_dense2_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [128, 128],
            'conv_filters': [32, 32, 32, 32, 32, 32],
            'conv_sizes': [5, 5, 5, 3, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2],

        },
    },
    'cnn6_dense3_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2],

        },
    },
    'cnn7_dense2_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 128],
            'conv_filters': [128, 128, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn8_dense2_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn9_dense2_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn9_dense3_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
    'cnn10_dense3_notrafo_flattened_pytorch': {
        'class': ConvNet,
        'framework': 'pytorch',
        'preprocessing': 'zstd-mono',
        'flatten_channels': True,
        'batch_size': 64,
        'early_stopping': True,
        'init_kwargs': {
            'output_units': '$n_classes',
            'dense_units': [256, 256, 128],
            'conv_filters': [128, 128, 128, 256, 256, 256, 256, 256, 256, 256],
            'conv_sizes': [9, 9, 9, 5, 5, 5, 5, 3, 3, 3],
            'conv_padding': 'same',
            'conv_padding_mode': 'zeros',
            'conv_strides': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            'pool_sizes': [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
            'pool_strides': [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],

        },
    },
}
