def mypad(x, pad, mode='constant', value=0): if mode == 'symmetric': # Horizontal only if pad[0] == 0 and pad[1] == 0: m1, m2 = pad[2], pad[3] l = x.shape[-2] xe = reflect(np.arange(-m1, l + m2, dtype='int32'), -0.5, l - 0.5) return x[:, :, xe] # Vertical only elif pad[2] == 0 and pad[3] == 0: m1, m2 = pad[0], pad[1] l = x.shape[-1] xe = reflect(np.arange(-m1, l + m2, dtype='int32'), -0.5, l - 0.5) return x[:, :, :, xe] # Both else: m1, m2 = pad[0], pad[1] l1 = x.shape[-1] xe_row = reflect(np.arange(-m1, l1 + m2, dtype='int32'), -0.5, l1 - 0.5) m1, m2 = pad[2], pad[3] l2 = x.shape[-2] xe_col = reflect(np.arange(-m1, l2 + m2, dtype='int32'), -0.5, l2 - 0.5) i = np.outer(xe_col, np.ones(xe_row.shape[0])) j = np.outer(np.ones(xe_col.shape[0]), xe_row) return x[:, :, i, j] elif mode == 'constant' or mode == 'reflect' or mode == 'replicate': return F.pad(x, pad, mode, value) elif mode == 'zero': return F.pad(x, pad) else: raise ValueError("Unkown pad type: {}".format(mode))
def mypad(x, pad, mode='constant', value=0): """ Function to do numpy like padding on tensors. Only works for 2-D padding. Inputs: x (tensor): tensor to pad pad (tuple): tuple of (left, right, top, bottom) pad sizes mode (str): 'symmetric', 'wrap', 'constant, 'reflect', 'replicate', or 'zero'. The padding technique. """ if mode == 'symmetric': # Vertical only if pad[0] == 0 and pad[1] == 0: m1, m2 = pad[2], pad[3] l = x.shape[-2] xe = reflect(np.arange(-m1, l+m2, dtype='int32'), -0.5, l-0.5) return x[:,:,xe] # horizontal only elif pad[2] == 0 and pad[3] == 0: m1, m2 = pad[0], pad[1] l = x.shape[-1] xe = reflect(np.arange(-m1, l+m2, dtype='int32'), -0.5, l-0.5) return x[:,:,:,xe] # Both else: m1, m2 = pad[0], pad[1] l1 = x.shape[-1] xe_row = reflect(np.arange(-m1, l1+m2, dtype='int32'), -0.5, l1-0.5) m1, m2 = pad[2], pad[3] l2 = x.shape[-2] xe_col = reflect(np.arange(-m1, l2+m2, dtype='int32'), -0.5, l2-0.5) i = np.outer(xe_col, np.ones(xe_row.shape[0])) j = np.outer(np.ones(xe_col.shape[0]), xe_row) return x[:,:,i,j] elif mode == 'periodic': # Vertical only if pad[0] == 0 and pad[1] == 0: xe = np.arange(x.shape[-2]) xe = np.pad(xe, (pad[2], pad[3]), mode='wrap') return x[:,:,xe] # Horizontal only elif pad[2] == 0 and pad[3] == 0: xe = np.arange(x.shape[-1]) xe = np.pad(xe, (pad[0], pad[1]), mode='wrap') return x[:,:,:,xe] # Both else: xe_col = np.arange(x.shape[-2]) xe_col = np.pad(xe_col, (pad[2], pad[3]), mode='wrap') xe_row = np.arange(x.shape[-1]) xe_row = np.pad(xe_row, (pad[0], pad[1]), mode='wrap') i = np.outer(xe_col, np.ones(xe_row.shape[0])) j = np.outer(np.ones(xe_col.shape[0]), xe_row) return x[:,:,i,j] elif mode == 'constant' or mode == 'reflect' or mode == 'replicate': return F.pad(x, pad, mode, value) elif mode == 'zero': return F.pad(x, pad) else: raise ValueError("Unkown pad type: {}".format(mode))