HE/parser.py at master · morgan76/HE · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
# -*- coding: utf-8 -*-
import argparse
from pathlib import Path


def get_parser():
    parser = argparse.ArgumentParser(description='Neural Features trainer')

    # Dataset parameters
    parser.add_argument('model_name', type=str,
                        help='Name of the model. If it exists, resume'
                        ' training.')
    parser.add_argument('ds_path', type=str,
                        help='Path to the  dataset.')
    parser.add_argument('--bnd-shift', type=float)
    parser.add_argument('--exclude-silence', action='store_false',
                        help='Should silent segments be excluded?',
                        default=None)

    # Files parameters
    parser.add_argument('--checkpoints_dir', type=Path,
                        help="Folder where to store checkpoints.")
    parser.add_argument('--models_dir', type=Path,
                        help="Folers where to store trained models.")

    # Model parameters
    # parser.add_argument('--use-batch-norm', action='store_true')  #TODO: FIX that in parser
    # TODO: add kernel size in parser

    # Feature parameters
    parser.add_argument('--feat-id', type=str,
                        help='Name of the feature to train embedding on')
    parser.add_argument('--feat-type', type=str,
                        help='Name of the method used to compute features')
    parser.add_argument('--sr', type=int, help="Sample rate of input audio.")
    parser.add_argument('--n-fft', type=int)
    parser.add_argument('--hop-size', type=int)
    parser.add_argument('--n-embedding', type=int,
                        help="Number of feature frames in the embedding.")

    # Model parameters
    parser.add_argument('--nb-workers', type=int,
                        help='Number of workers for dataloader.')
    parser.add_argument('--margin', type=float,
                        help='margin parameter of the triplet loss function')
    parser.add_argument('-p', type=int,
                        help='Norm degree for pairwise distance of triplet'
                        ' loss')

    # Training parameters
    parser.add_argument('-e', '--epochs', type=int)
    parser.add_argument('-b', '--batch-size', type=int)
    parser.add_argument('--lr', type=float, dest='learning_rate',
                        help='learning rate, defaults to 1e-2')
    parser.add_argument('--mining', type=str,
                        help='Strategy for triplet mining')
    parser.add_argument('--m-per-class', type=int,
                        help="Number of samples per class in each batch")
    parser.add_argument('--classes', type=str, dest='classes_strategy',
                        choices=('track', 'label'))
    parser.add_argument('--seed', type=int, metavar='S',
                        help='random seed (default: 42)')
    parser.add_argument('--no-cuda', action='store_true', default=None,
                        help='disables CUDA training')
    parser.add_argument('--quiet', action='store_true', default=None,
                        help='less verbose during training')
    parser.add_argument('--resume', action='store_true', default=False,
                        help="Parameter to resume training of a model.")

    return parser