-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathtest.py
178 lines (139 loc) · 7.08 KB
/
test.py
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
174
175
176
177
178
"""
This module containg all the necessary methods for testing the models trained to predict values
for valence and arousal.
"""
import os
import torch
from models import AudioNet
from data_loader import make_testing_loader
from utility_functions import *
class Tester:
"""
Methods for testing are defined in this class.
Attributes:
dimension (str): specifies the type of output predicted by the model
test_loader: loading and batching the data in test set
model: AudioNet model with parameters according to `dimension`
valence_dict, arousal_dict, quadrants_dict: dictionaries containing data needed for computing
performance metrics and visualization
"""
def __init__(self, args):
self.dimension = args.dimension
self._data_dir = args.data_dir
self._models_dir = args.models_dir
self._plots_dir = args.plots_dir
self._device = torch.device(args.device if torch.cuda.is_available() else 'cpu')
self.test_loader = make_testing_loader(self._data_dir)
if self.dimension == 'both':
self._params_dict = args.params_dict
self.model = AudioNet(self._params_dict).to(self._device)
else:
self._valence_params_dict = args.valence_params_dict
self._arousal_params_dict = args.arousal_params_dict
self.valence_model = AudioNet(self._valence_params_dict).to(self._device)
self.arousal_model = AudioNet(self._arousal_params_dict).to(self._device)
self.valence_dict = dict()
self.arousal_dict = dict()
self.quadrants_dict = dict()
def load_model_1d(self):
"""
Method to load the pretrained models to predict separately values for valence and arousal dimension,
respectively.
"""
valence_path = os.path.join(self._models_dir, 'model_valence.pt')
self.valence_model.load_state_dict(torch.load(valence_path))
arousal_path = os.path.join(self._models_dir, 'model_arousal.pt')
self.arousal_model.load_state_dict(torch.load(arousal_path))
def load_model_2d(self):
"""
Method to load the pretrained model to predict values for both valence and arousal dimensions.
"""
model_path = os.path.join(self._models_dir, 'model_{:s}.pt'.format(self.dimension))
self.model.load_state_dict(torch.load(model_path))
def test_1d(self):
"""
Method to test 1D-output models. This is called when `dimension` is `valence` or `arousal`.
"""
true_valence = []
pred_valence = []
true_arousal = []
pred_arousal = []
self.valence_model.eval()
self.arousal_model.eval()
# Freeze gradients
with torch.no_grad():
for batch_idx, (data, annotations) in enumerate(self.test_loader):
# Create individual target labels for each dimension
valence_target = annotations[:, 0]
arousal_target = annotations[:, 1]
# Move data to device
data = data.to(self._device)
valence_target = valence_target.to(self._device)
arousal_target = arousal_target.to(self._device)
# Make predictions for valence
valence_output = self.valence_model(data)
valence_target = valence_target.view_as(valence_output)
# Make predictions for arousal
arousal_output = self.arousal_model(data)
arousal_target = arousal_target.view_as(arousal_output)
true_valence.extend(valence_target.cpu().detach().squeeze().numpy())
pred_valence.extend(valence_output.cpu().detach().squeeze().numpy())
true_arousal.extend(arousal_target.cpu().detach().squeeze().numpy())
pred_arousal.extend(arousal_output.cpu().detach().squeeze().numpy())
true_valence, pred_valence = np.array(true_valence), np.array(pred_valence)
# Compute valence MAE and MSE
valence_mae = np.mean(np.abs(true_valence - pred_valence))
valence_mse = np.mean((true_valence - pred_valence) ** 2)
# Compute arousal MAE and MSE
true_arousal, pred_arousal = np.array(true_arousal), np.array(pred_arousal)
arousal_mae = np.mean(np.abs(true_arousal - pred_arousal))
arousal_mse = np.mean((true_arousal - pred_arousal) ** 2)
self.valence_dict['true_annotations'] = true_valence
self.valence_dict['pred_annotations'] = pred_valence
self.valence_dict['mae'] = valence_mae
self.valence_dict['mse'] = valence_mse
self.arousal_dict['true_annotations'] = true_arousal
self.arousal_dict['pred_annotations'] = pred_arousal
self.arousal_dict['mae'] = arousal_mae
self.arousal_dict['mse'] = arousal_mse
# Extract information about quadrants from valence & arousal annotations
self.quadrants_dict = get_quadrants_dict(self.valence_dict, self.arousal_dict)
def test_2d(self):
"""
Method to test 2D-output models. This is called when `dimension` is `both`.
"""
true_annotations = []
pred_annotations = []
self.model.eval()
# Freeze gradients
with torch.no_grad():
for batch_idx, (data, annotations) in enumerate(self.test_loader):
# Move data to device
data = data.to(self._device)
annotations = annotations.to(self._device)
# Make predictions for valence and arousal
output = self.model(data)
true_annotations.extend(annotations.cpu().detach().numpy())
pred_annotations.extend(output.cpu().detach().numpy())
true_annotations = np.array(true_annotations)
pred_annotations = np.array(pred_annotations)
# Compute valence MAE and MSE
true_valence = np.array([annot[0] for annot in true_annotations])
pred_valence = np.array([annot[0] for annot in pred_annotations])
valence_mae = np.mean(np.abs(true_valence - pred_valence))
valence_mse = np.mean((true_valence - pred_valence) ** 2)
# Extract predictions and true values for valence and arousal dimensions and compute MAE and MSE
true_arousal = np.array([annot[1] for annot in true_annotations])
pred_arousal = np.array([annot[1] for annot in pred_annotations])
arousal_mae = np.mean(np.abs(true_arousal - pred_arousal))
arousal_mse = np.mean((true_arousal - pred_arousal) ** 2)
self.valence_dict['true_annotations'] = true_valence
self.valence_dict['pred_annotations'] = pred_valence
self.valence_dict['mae'] = valence_mae
self.valence_dict['mse'] = valence_mse
self.arousal_dict['true_annotations'] = true_arousal
self.arousal_dict['pred_annotations'] = pred_arousal
self.arousal_dict['mae'] = arousal_mae
self.arousal_dict['mse'] = arousal_mse
# Extract information about quadrants from valence & arousal annotations
self.quadrants_dict = get_quadrants_dict(self.valence_dict, self.arousal_dict)