Source code for espnet.utils.spec_augment

# -*- coding: utf-8 -*-

This implementation is modified from

MIT License

Copyright (c) 2019 Zach Caceres

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.


import random

import torch

[docs]def specaug( spec, W=5, F=30, T=40, num_freq_masks=2, num_time_masks=2, replace_with_zero=False ): """SpecAugment Reference: SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition ( This implementation modified from :param torch.Tensor spec: input tensor with the shape (T, dim) :param int W: time warp parameter :param int F: maximum width of each freq mask :param int T: maximum width of each time mask :param int num_freq_masks: number of frequency masks :param int num_time_masks: number of time masks :param bool replace_with_zero: if True, masked parts will be filled with 0, if False, filled with mean """ return time_mask( freq_mask( time_warp(spec, W=W), F=F, num_masks=num_freq_masks, replace_with_zero=replace_with_zero, ), T=T, num_masks=num_time_masks, replace_with_zero=replace_with_zero, )
[docs]def time_warp(spec, W=5): """Time warping :param torch.Tensor spec: input tensor with shape (T, dim) :param int W: time warp parameter """ spec = spec.unsqueeze(0) spec_len = spec.shape[1] num_rows = spec.shape[2] device = spec.device y = num_rows // 2 horizontal_line_at_ctr = spec[0, :, y] assert len(horizontal_line_at_ctr) == spec_len point_to_warp = horizontal_line_at_ctr[random.randrange(W, spec_len - W)] assert isinstance(point_to_warp, torch.Tensor) # Uniform distribution from (0,W) with chance to be up to W negative dist_to_warp = random.randrange(-W, W) src_pts, dest_pts = ( torch.tensor([[[point_to_warp, y]]], device=device), torch.tensor([[[point_to_warp + dist_to_warp, y]]], device=device), ) warped_spectro, dense_flows = sparse_image_warp(spec, src_pts, dest_pts) return warped_spectro.squeeze(3).squeeze(0)
[docs]def freq_mask(spec, F=30, num_masks=1, replace_with_zero=False): """Frequency masking :param torch.Tensor spec: input tensor with shape (T, dim) :param int F: maximum width of each mask :param int num_masks: number of masks :param bool replace_with_zero: if True, masked parts will be filled with 0, if False, filled with mean """ cloned = spec.unsqueeze(0).clone() num_mel_channels = cloned.shape[2] for i in range(0, num_masks): f = random.randrange(0, F) f_zero = random.randrange(0, num_mel_channels - f) # avoids randrange error if values are equal and range is empty if f_zero == f_zero + f: return cloned.squeeze(0) mask_end = random.randrange(f_zero, f_zero + f) if replace_with_zero: cloned[0][:, f_zero:mask_end] = 0 else: cloned[0][:, f_zero:mask_end] = cloned.mean() return cloned.squeeze(0)
[docs]def time_mask(spec, T=40, num_masks=1, replace_with_zero=False): """Time masking :param torch.Tensor spec: input tensor with shape (T, dim) :param int T: maximum width of each mask :param int num_masks: number of masks :param bool replace_with_zero: if True, masked parts will be filled with 0, if False, filled with mean """ cloned = spec.unsqueeze(0).clone() len_spectro = cloned.shape[1] for i in range(0, num_masks): t = random.randrange(0, T) t_zero = random.randrange(0, len_spectro - t) # avoids randrange error if values are equal and range is empty if t_zero == t_zero + t: return cloned.squeeze(0) mask_end = random.randrange(t_zero, t_zero + t) if replace_with_zero: cloned[0][t_zero:mask_end, :] = 0 else: cloned[0][t_zero:mask_end, :] = cloned.mean() return cloned.squeeze(0)
[docs]def sparse_image_warp( img_tensor, source_control_point_locations, dest_control_point_locations, interpolation_order=2, regularization_weight=0.0, num_boundaries_points=0, ): device = img_tensor.device control_point_flows = dest_control_point_locations - source_control_point_locations batch_size, image_height, image_width = img_tensor.shape flattened_grid_locations = get_flat_grid_locations( image_height, image_width, device ) flattened_flows = interpolate_spline( dest_control_point_locations, control_point_flows, flattened_grid_locations, interpolation_order, regularization_weight, ) dense_flows = create_dense_flows( flattened_flows, batch_size, image_height, image_width ) warped_image = dense_image_warp(img_tensor, dense_flows) return warped_image, dense_flows
[docs]def get_grid_locations(image_height, image_width, device): y_range = torch.linspace(0, image_height - 1, image_height, device=device) x_range = torch.linspace(0, image_width - 1, image_width, device=device) y_grid, x_grid = torch.meshgrid(y_range, x_range) return torch.stack((y_grid, x_grid), -1)
[docs]def flatten_grid_locations(grid_locations, image_height, image_width): return torch.reshape(grid_locations, [image_height * image_width, 2])
[docs]def get_flat_grid_locations(image_height, image_width, device): y_range = torch.linspace(0, image_height - 1, image_height, device=device) x_range = torch.linspace(0, image_width - 1, image_width, device=device) y_grid, x_grid = torch.meshgrid(y_range, x_range) return torch.stack((y_grid, x_grid), -1).reshape([image_height * image_width, 2])
[docs]def create_dense_flows(flattened_flows, batch_size, image_height, image_width): # possibly .view return torch.reshape(flattened_flows, [batch_size, image_height, image_width, 2])
[docs]def interpolate_spline( train_points, train_values, query_points, order, regularization_weight=0.0, ): # First, fit the spline to the observed data. w, v = solve_interpolation(train_points, train_values, order, regularization_weight) # Then, evaluate the spline at the query locations. query_values = apply_interpolation(query_points, train_points, w, v, order) return query_values
[docs]def solve_interpolation(train_points, train_values, order, regularization_weight): device = train_points.device b, n, d = train_points.shape k = train_values.shape[-1] c = train_points f = train_values.float() matrix_a = phi(cross_squared_distance_matrix(c, c), order).unsqueeze(0) # [b, n, n] # Append ones to the feature values for the bias term in the linear model. ones = torch.ones(1, dtype=train_points.dtype, device=device).view([-1, 1, 1]) matrix_b =, ones), 2).float() # [b, n, d + 1] # [b, n + d + 1, n] left_block =, torch.transpose(matrix_b, 2, 1)), 1) num_b_cols = matrix_b.shape[2] # d + 1 # In Tensorflow, zeros are used here. Pytorch solve fails with zeros # for some reason we don't understand. # So instead we use very tiny randn values (variance of one, zero mean) # on one side of our multiplication. lhs_zeros = torch.randn((b, num_b_cols, num_b_cols), device=device) / 1e10 right_block =, lhs_zeros), 1) # [b, n + d + 1, d + 1] lhs =, right_block), 2) # [b, n + d + 1, n + d + 1] rhs_zeros = torch.zeros( (b, d + 1, k), dtype=train_points.dtype, device=device ).float() rhs =, rhs_zeros), 1) # [b, n + d + 1, k] # Then, solve the linear system and unpack the results. X, LU = torch.gesv(rhs, lhs) w = X[:, :n, :] v = X[:, n:, :] return w, v
[docs]def cross_squared_distance_matrix(x, y): """Pairwise squared distance between two (batch) matrices' rows (2nd dim). Computes the pairwise distances between rows of x and rows of y Args: x: [batch_size, n, d] float `Tensor` y: [batch_size, m, d] float `Tensor` Returns: squared_dists: [batch_size, n, m] float `Tensor`, where squared_dists[b,i,j] = ||x[b,i,:] - y[b,j,:]||^2 """ x_norm_squared = torch.sum(torch.mul(x, x)) y_norm_squared = torch.sum(torch.mul(y, y)) x_y_transpose = torch.matmul(x.squeeze(0), y.squeeze(0).transpose(0, 1)) # squared_dists[b,i,j] = ||x_bi - y_bj||^2 = x_bi'x_bi- 2x_bi'x_bj + x_bj'x_bj squared_dists = x_norm_squared - 2 * x_y_transpose + y_norm_squared return squared_dists.float()
[docs]def phi(r, order): """Coordinate-wise nonlinearity used to define the order of the interpolation. See for the definition. Args: r: input op order: interpolation order Returns: phi_k evaluated coordinate-wise on r, for k = r """ EPSILON = torch.tensor(1e-10, device=r.device) # using EPSILON prevents log(0), sqrt0), etc. # sqrt(0) is well-defined, but its gradient is not if order == 1: r = torch.max(r, EPSILON) r = torch.sqrt(r) return r elif order == 2: return 0.5 * r * torch.log(torch.max(r, EPSILON)) elif order == 4: return 0.5 * torch.square(r) * torch.log(torch.max(r, EPSILON)) elif order % 2 == 0: r = torch.max(r, EPSILON) return 0.5 * torch.pow(r, 0.5 * order) * torch.log(r) else: r = torch.max(r, EPSILON) return torch.pow(r, 0.5 * order)
[docs]def apply_interpolation(query_points, train_points, w, v, order): """Apply polyharmonic interpolation model to data. Notes: Given coefficients w and v for the interpolation model, we evaluate interpolated function values at query_points. Args: query_points: `[b, m, d]` x values to evaluate the interpolation at train_points: `[b, n, d]` x values that act as the interpolation centers ( the c variables in the wikipedia article) w: `[b, n, k]` weights on each interpolation center v: `[b, d, k]` weights on each input dimension order: order of the interpolation Returns: Polyharmonic interpolation evaluated at points defined in query_points. """ query_points = query_points.unsqueeze(0) # First, compute the contribution from the rbf term. pairwise_dists = cross_squared_distance_matrix( query_points.float(), train_points.float() ) phi_pairwise_dists = phi(pairwise_dists, order) rbf_term = torch.matmul(phi_pairwise_dists, w) # Then, compute the contribution from the linear term. # Pad query_points with ones, for the bias term in the linear model. ones = torch.ones_like(query_points[..., :1]) query_points_pad =, ones), 2).float() linear_term = torch.matmul(query_points_pad, v) return rbf_term + linear_term
[docs]def dense_image_warp(image, flow): """Image warping using per-pixel flow vectors. Apply a non-linear warp to the image, where the warp is specified by a dense flow field of offset vectors that define the correspondences of pixel values in the output image back to locations in the source image. Specifically, the pixel value at output[b, j, i, c] is images[b, j - flow[b, j, i, 0], i - flow[b, j, i, 1], c]. The locations specified by this formula do not necessarily map to an int index. Therefore, the pixel value is obtained by bilinear interpolation of the 4 nearest pixels around (b, j - flow[b, j, i, 0], i - flow[b, j, i, 1]). For locations outside of the image, we use the nearest pixel values at the image boundary. Args: image: 4-D float `Tensor` with shape `[batch, height, width, channels]`. flow: A 4-D float `Tensor` with shape `[batch, height, width, 2]`. name: A name for the operation (optional). Note that image and flow can be of type tf.half, tf.float32, or tf.float64, and do not necessarily have to be the same type. Returns: A 4-D float `Tensor` with shape`[batch, height, width, channels]` and same type as input image. Raises: ValueError: if height < 2 or width < 2 or the inputs have the wrong number of dimensions. """ image = image.unsqueeze(3) # add a single channel dimension to image tensor batch_size, height, width, channels = image.shape device = image.device # The flow is defined on the image grid. Turn the flow into a list of query # points in the grid space. grid_x, grid_y = torch.meshgrid( torch.arange(width, device=device), torch.arange(height, device=device) ) stacked_grid = torch.stack((grid_y, grid_x), dim=2).float() batched_grid = stacked_grid.unsqueeze(-1).permute(3, 1, 0, 2) query_points_on_grid = batched_grid - flow query_points_flattened = torch.reshape( query_points_on_grid, [batch_size, height * width, 2] ) # Compute values at the query points, then reshape the result back to the # image grid. interpolated = interpolate_bilinear(image, query_points_flattened) interpolated = torch.reshape(interpolated, [batch_size, height, width, channels]) return interpolated
[docs]def interpolate_bilinear( grid, query_points, name="interpolate_bilinear", indexing="ij" ): """Similar to Matlab's interp2 function. Notes: Finds values for query points on a grid using bilinear interpolation. Args: grid: a 4-D float `Tensor` of shape `[batch, height, width, channels]`. query_points: a 3-D float `Tensor` of N points with shape `[batch, N, 2]`. name: a name for the operation (optional). indexing: whether the query points are specified as row and column (ij), or Cartesian coordinates (xy). Returns: values: a 3-D `Tensor` with shape `[batch, N, channels]` Raises: ValueError: if the indexing mode is invalid, or if the shape of the inputs invalid. """ if indexing != "ij" and indexing != "xy": raise ValueError("Indexing mode must be 'ij' or 'xy'") shape = grid.shape if len(shape) != 4: msg = "Grid must be 4 dimensional. Received size: " raise ValueError(msg + str(grid.shape)) batch_size, height, width, channels = grid.shape shape = [batch_size, height, width, channels] query_type = query_points.dtype grid_type = grid.dtype grid_device = grid.device num_queries = query_points.shape[1] alphas = [] floors = [] ceils = [] index_order = [0, 1] if indexing == "ij" else [1, 0] unstacked_query_points = query_points.unbind(2) for dim in index_order: queries = unstacked_query_points[dim] size_in_indexing_dimension = shape[dim + 1] # max_floor is size_in_indexing_dimension - 2 so that max_floor + 1 # is still a valid index into the grid. max_floor = torch.tensor( size_in_indexing_dimension - 2, dtype=query_type, device=grid_device ) min_floor = torch.tensor(0.0, dtype=query_type, device=grid_device) maxx = torch.max(min_floor, torch.floor(queries)) floor = torch.min(maxx, max_floor) int_floor = floor.long() floors.append(int_floor) ceil = int_floor + 1 ceils.append(ceil) # alpha has the same type as the grid, as we will directly use alpha # when taking linear combinations of pixel values from the image. alpha = torch.tensor((queries - floor), dtype=grid_type, device=grid_device) min_alpha = torch.tensor(0.0, dtype=grid_type, device=grid_device) max_alpha = torch.tensor(1.0, dtype=grid_type, device=grid_device) alpha = torch.min(torch.max(min_alpha, alpha), max_alpha) # Expand alpha to [b, n, 1] so we can use broadcasting # (since the alpha values don't depend on the channel). alpha = torch.unsqueeze(alpha, 2) alphas.append(alpha) flattened_grid = torch.reshape(grid, [batch_size * height * width, channels]) batch_offsets = torch.reshape( torch.arange(batch_size, device=grid_device) * height * width, [batch_size, 1] ) # This wraps array_ops.gather. We reshape the image data such that the # batch, y, and x coordinates are pulled into the first dimension. # Then we gather. Finally, we reshape the output back. It's possible this # code would be made simpler by using array_ops.gather_nd. def gather(y_coords, x_coords, name): linear_coordinates = batch_offsets + y_coords * width + x_coords gathered_values = torch.gather(flattened_grid.t(), 1, linear_coordinates) return torch.reshape(gathered_values, [batch_size, num_queries, channels]) # grab the pixel values in the 4 corners around each query point top_left = gather(floors[0], floors[1], "top_left") top_right = gather(floors[0], ceils[1], "top_right") bottom_left = gather(ceils[0], floors[1], "bottom_left") bottom_right = gather(ceils[0], ceils[1], "bottom_right") interp_top = alphas[1] * (top_right - top_left) + top_left interp_bottom = alphas[1] * (bottom_right - bottom_left) + bottom_left interp = alphas[0] * (interp_bottom - interp_top) + interp_top return interp