Alex Blank
100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
import numpy as np
|
|
import torch
|
|
|
|
|
|
def flat_x_y_collate(batch, *args, **kwargs) -> tuple:
|
|
"""
|
|
Collate function for a classical model with flat x and y feature vectors
|
|
"""
|
|
# Flatten each item in the batch
|
|
flattened_x = list()
|
|
flattened_y = list()
|
|
for i in range(len(batch)):
|
|
item = batch[i]
|
|
flattened_x_item = list()
|
|
flattened_y_item = list()
|
|
for j, (category, windows) in enumerate(item.items()):
|
|
for feature in windows.T:
|
|
if category == "target":
|
|
flattened_y_item.append(feature)
|
|
else:
|
|
flattened_x_item.append(feature)
|
|
# create numpy arrays with t_1_f1, t_1_f2, t_2_f1, t_2_f2 and so on
|
|
reordered_x = np.empty((len(flattened_x_item) * len(flattened_x_item[0])))
|
|
reordered_y = np.empty((len(flattened_y_item) * len(flattened_y_item[0])))
|
|
for j in range(len(flattened_x_item)):
|
|
for k in range(len(flattened_x_item[j])):
|
|
reordered_x[j + k * len(flattened_x_item)] = flattened_x_item[j][k]
|
|
|
|
for j in range(len(flattened_y_item)):
|
|
for k in range(len(flattened_y_item[j])):
|
|
reordered_y[j + k * len(flattened_y_item)] = flattened_y_item[j][k]
|
|
|
|
# append to the list
|
|
flattened_x.append(reordered_x)
|
|
flattened_y.append(reordered_y)
|
|
|
|
return np.array(flattened_x), np.array(flattened_y)
|
|
|
|
|
|
def simple_x_y_collate(batch):
|
|
"""
|
|
Collate function for LSTM model
|
|
"""
|
|
collated_x = list()
|
|
collated_y = list()
|
|
for item in batch:
|
|
current_x = list()
|
|
current_y = list()
|
|
for category, windows in item.items():
|
|
if category == "target_features":
|
|
current_y.append(windows)
|
|
else:
|
|
current_x.append(windows)
|
|
|
|
collated_x.append(np.concatenate(current_x, axis=1))
|
|
collated_y.append(np.concatenate(current_y, axis=1))
|
|
|
|
# convert to numpy arrays
|
|
collated_x = np.array(collated_x)
|
|
collated_y = np.array(collated_y)
|
|
|
|
# convert to torch tensors
|
|
return torch.tensor(collated_x, dtype=torch.float32), torch.tensor(collated_y, dtype=torch.float32)
|
|
|
|
|
|
def collate_with_padding(batch,
|
|
padding_value: float = 0.0, ):
|
|
"""
|
|
Collate the batch with padding
|
|
"""
|
|
|
|
# create simple lists for x and y
|
|
collated_x = list()
|
|
collated_y = list()
|
|
for item in batch:
|
|
current_x = list()
|
|
current_y = list()
|
|
for category, windows in item.items():
|
|
if category == "target_features":
|
|
current_y.append(windows)
|
|
else:
|
|
current_x.append(windows)
|
|
|
|
collated_x.append(np.concatenate(current_x, axis=1))
|
|
collated_y.append(np.concatenate(current_y, axis=1))
|
|
|
|
# get the max length of the x and y
|
|
max_x_length = max([x.shape[0] for x in collated_x])
|
|
max_y_length = max([y.shape[0] for y in collated_y])
|
|
# pad the x and y
|
|
padded_x = list()
|
|
padded_y = list()
|
|
for x, y in zip(collated_x, collated_y):
|
|
padded_x.append(
|
|
np.pad(x, ((max_x_length - x.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value))
|
|
padded_y.append(
|
|
np.pad(y, ((max_y_length - y.shape[0], 0), (0, 0)), mode='constant', constant_values=padding_value))
|
|
|
|
return torch.tensor(padded_x, dtype=torch.float32), torch.tensor(padded_y, dtype=torch.float32)
|