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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
|
"""
Training the neural pitch estimator
"""
import os
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('features', type=str, help='.f32 IF Features for training (generated by augmentation script)')
parser.add_argument('features_pitch', type=str, help='.npy Pitch file for training (generated by augmentation script)')
parser.add_argument('output_folder', type=str, help='Output directory to store the model weights and config')
parser.add_argument('data_format', type=str, help='Choice of Input Data',choices=['if','xcorr','both'])
parser.add_argument('--gpu_index', type=int, help='GPU index to use if multiple GPUs',default = 0,required = False)
parser.add_argument('--confidence_threshold', type=float, help='Confidence value below which pitch will be neglected during training',default = 0.4,required = False)
parser.add_argument('--context', type=int, help='Sequence length during training',default = 100,required = False)
parser.add_argument('--N', type=int, help='STFT window size',default = 320,required = False)
parser.add_argument('--H', type=int, help='STFT Hop size',default = 160,required = False)
parser.add_argument('--xcorr_dimension', type=int, help='Dimension of Input cross-correlation',default = 257,required = False)
parser.add_argument('--freq_keep', type=int, help='Number of Frequencies to keep',default = 30,required = False)
parser.add_argument('--gru_dim', type=int, help='GRU Dimension',default = 64,required = False)
parser.add_argument('--output_dim', type=int, help='Output dimension',default = 192,required = False)
parser.add_argument('--learning_rate', type=float, help='Learning Rate',default = 1.0e-3,required = False)
parser.add_argument('--epochs', type=int, help='Number of training epochs',default = 50,required = False)
parser.add_argument('--choice_cel', type=str, help='Choice of Cross Entropy Loss (default or robust)',choices=['default','robust'],default = 'default',required = False)
parser.add_argument('--prefix', type=str, help="prefix for model export, default: model", default='model')
args = parser.parse_args()
# import os
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
# os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu_index)
# Fixing the seeds for reproducability
import time
np_seed = int(time.time())
torch_seed = int(time.time())
import json
import torch
torch.manual_seed(torch_seed)
import numpy as np
np.random.seed(np_seed)
from utils import count_parameters
import tqdm
import sys
from datetime import datetime
#from evaluation import rpa
# print(list(range(torch.cuda.device_count())))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# device = 'cpu'
from models import loader_joint as loader
if args.data_format == 'if':
from models import large_if_ccode as model
pitch_nn = model(args.freq_keep*3,args.gru_dim,args.output_dim)
elif args.data_format == 'xcorr':
from models import large_xcorr as model
pitch_nn = model(args.xcorr_dimension,args.gru_dim,args.output_dim)
else:
from models import large_joint as model
pitch_nn = model(88,224,args.gru_dim,args.output_dim)
dataset_training = loader(args.features,args.features_pitch,args.confidence_threshold,args.context,args.data_format)
def loss_custom(logits,labels,confidence,choice = 'default',nmax = 192,q = 0.7):
logits_softmax = torch.nn.Softmax(dim = 1)(logits).permute(0,2,1)
labels_one_hot = torch.nn.functional.one_hot(labels.long(),nmax)
if choice == 'default':
# Categorical Cross Entropy
CE = -torch.sum(torch.log(logits_softmax*labels_one_hot + 1.0e-6)*labels_one_hot,dim=-1)
CE = torch.mean(confidence*CE)
else:
# Robust Cross Entropy
CE = (1.0/q)*(1 - torch.sum(torch.pow(logits_softmax*labels_one_hot + 1.0e-7,q),dim=-1) )
CE = torch.sum(confidence*CE)
return CE
def accuracy(logits,labels,confidence,choice = 'default',nmax = 192,q = 0.7):
logits_softmax = torch.nn.Softmax(dim = 1)(logits).permute(0,2,1)
pred_pitch = torch.argmax(logits_softmax, 2)
#print(pred_pitch.shape, labels.long().shape)
accuracy = (pred_pitch != labels.long())*1.
#print(accuracy.shape, confidence.shape)
return 1.-torch.mean(confidence*accuracy)
# features = args.features
# pitch = args.crepe_pitch
# dataset_training = loader(features,pitch,args.confidence_threshold,args.freq_keep,args.context)
# dataset_training = loader(features,pitch,'../../../../testing/testing_features_10pct_xcorr.f32')
train_dataset, test_dataset = torch.utils.data.random_split(dataset_training, [0.95,0.05],generator=torch.Generator().manual_seed(torch_seed))
batch_size = 256
train_dataloader = torch.utils.data.DataLoader(dataset = train_dataset,batch_size = batch_size,shuffle = True,num_workers = 0, pin_memory = False)
test_dataloader = torch.utils.data.DataLoader(dataset = test_dataset,batch_size = batch_size,shuffle = True,num_workers = 0, pin_memory = False)
# pitch_nn = model(args.freq_keep*3,args.gru_dim,args.output_dim).to(device)
pitch_nn = pitch_nn.to(device)
num_params = count_parameters(pitch_nn)
learning_rate = args.learning_rate
model_opt = torch.optim.Adam(pitch_nn.parameters(), lr = learning_rate)
num_epochs = args.epochs
for epoch in range(num_epochs):
losses = []
accs = []
pitch_nn.train()
with tqdm.tqdm(train_dataloader) as train_epoch:
for i, (xi, yi, ci) in enumerate(train_epoch):
yi, xi, ci = yi.to(device, non_blocking=True), xi.to(device, non_blocking=True), ci.to(device, non_blocking=True)
pi = pitch_nn(xi.float())
loss = loss_custom(logits = pi,labels = yi,confidence = ci,choice = args.choice_cel,nmax = args.output_dim)
acc = accuracy(logits = pi,labels = yi,confidence = ci,choice = args.choice_cel,nmax = args.output_dim)
acc = acc.detach()
model_opt.zero_grad()
loss.backward()
model_opt.step()
losses.append(loss.item())
accs.append(acc.item())
avg_loss = np.mean(losses)
avg_acc = np.mean(accs)
train_epoch.set_postfix({"Train Epoch" : epoch, "Train Loss":avg_loss, "acc" : avg_acc.item()})
if epoch % 5 == 0:
pitch_nn.eval()
losses = []
with tqdm.tqdm(test_dataloader) as test_epoch:
for i, (xi, yi, ci) in enumerate(test_epoch):
yi, xi, ci = yi.to(device, non_blocking=True), xi.to(device, non_blocking=True), ci.to(device, non_blocking=True)
pi = pitch_nn(xi.float())
loss = loss_custom(logits = pi,labels = yi,confidence = ci,choice = args.choice_cel,nmax = args.output_dim)
losses.append(loss.item())
avg_loss = np.mean(losses)
test_epoch.set_postfix({"Epoch" : epoch, "Test Loss":avg_loss})
pitch_nn.eval()
#rpa(pitch_nn,device,data_format = args.data_format)
config = dict(
data_format = args.data_format,
epochs = num_epochs,
window_size = args.N,
hop_factor = args.H,
freq_keep = args.freq_keep,
batch_size = batch_size,
learning_rate = learning_rate,
confidence_threshold = args.confidence_threshold,
model_parameters = num_params,
np_seed = np_seed,
torch_seed = torch_seed,
xcorr_dim = args.xcorr_dimension,
dim_input = 3*args.freq_keep,
gru_dim = args.gru_dim,
output_dim = args.output_dim,
choice_cel = args.choice_cel,
context = args.context,
)
model_save_path = os.path.join(args.output, f"{args.prefix}_{args.data_format}.pth")
checkpoint = {
'state_dict': pitch_nn.state_dict(),
'config': config
}
torch.save(checkpoint, model_save_path)
|
