def __init__(self):
super(Net, self).__init__()
# self.lk1 = LongConv(10, 12288)
# self.lk2 = LongConv(10, 108300)
# self.lk3 = LongConv(100, 3072)
# self.lk4 = LongConv(50, 3072)
# self.lk5 = LongConv(50, 3072)
# self.lk6 = LongConv(50, 3072)
#
# self.lc = Looper(100, 3072)
# # self.lm = LongMem(1, 3072)
# self.lm = LongMem(10, 12288, 3, 3)
# self.lm2 = LongMem(10, 13872, 3, 3)
# self.lm3 = LongMem(10, 15552, 3, 3)
#
# self.conv = nn.Conv2d(3, 3, (3, 3), stride=1)
# self.conv2 = nn.Conv2d(3, 3, (3, 3), stride=1)
# self.conv3 = nn.Conv2d(3, 3, (3, 3), stride=1)
#
# self.tconv = nn.ConvTranspose2d(3, 3, (10, 10))
# self.tconv2 = nn.ConvTranspose2d(3, 3, (10, 10))
# self.m1 = MemA(10, 75*784, 75 * 15376)
# self.m2 = MemA(10, 75 * 784, 75 * 15376)
# self.m3 = MemA(10, 75 * 784, 75 * 15376)
# self.m4 = MemA(10, 32 * 32 * 3, 128 * 128 * 3)
# self.m5 = MemA(10, 16 * 16 * 3, 128 * 128 * 3)
self.mb1 = MemB(10, 128 * 128 * 3, 128 * 128 * 3)
# self.m2 = MemC(10, 32*32*3, 128*128*3)
self.unfold = nn.Unfold(kernel_size=(5, 5))
self.fold = nn.Fold(kernel_size=(5, 5),
output_size=(128, 128),
stride=5)
self.patches = nn.Parameter(torch.randn((1, 75, 15376)))
def __init__(self, kernel_size, output_size=None, dilation=1, stride=1, device=torch.device('cpu')):
super(AvgFeatAGG2d, self).__init__()
self.device = device
self.kernel_size = kernel_size
self.unfold = nn.Unfold(kernel_size=kernel_size, dilation=dilation, stride=stride)
self.fold = nn.Fold(output_size=output_size, kernel_size=1, dilation=1, stride=1)
self.output_size = output_size
def __init__(self,
input_size,
num_channels,
num_filters,
batch_size,
kernel_size,
learning_rate,
act_fn,
padding=0,
stride=1,
device="cpu"):
self.input_size = input_size
self.num_channels = num_channels
self.num_filters = num_filters
self.batch_size = batch_size
self.kernel_size = kernel_size
self.padding = padding
self.stride = stride
self.output_size = math.floor(
(self.input_size +
(2 * self.padding) - self.kernel_size) / self.stride) + 1
self.learning_rate = learning_rate
self.act_fn
self.device = device
self.kernel = torch.empty(self.num_filters, self.num_channels,
self.kernel_size, self.kernel_size).normal_(
mean=0, std=0.05).to(self.device)
self.unfold = nn.Unfold(kernel_size=(self.kernel_size,
self.kernel_size),
padding=self.padding,
stride=self.stride).to(self.device)
self.fold = nn.Fold(output_size=(self.input_size, self.input_size),
kernel_size=(self.kernel_size, self.kernel_size),
padding=self.padding,
stride=self.stride).to(self.device)
def image_from_patches(tensor_list,
mask_t,
base_tensor,
t_size,
tile_size=256,
padding=0):
stride = tile_size // 2
base_tensor = base_tensor.to('cuda')
# base_tensor here is used as a container
for t, tile in enumerate(tensor_list):
#print(tile.size())
#tile = torch.from_numpy(tile)
tile = tile.type(torch.FloatTensor)
tile = tile.to('cuda')
base_tensor[[t], :, :] = tile
base_tensor = base_tensor.permute(1, 2, 3, 0).reshape(
3 * tile_size * tile_size, base_tensor.size(0)).unsqueeze(0)
fold = nn.Fold(output_size=(t_size[0], t_size[1]),
kernel_size=(tile_size, tile_size),
stride=stride)
# https://discuss.pytorch.org/t/seemlessly-blending-tensors-together/65235/2?u=bowenroom
mask_t = mask_t.to('cuda')
output_tensor = fold(base_tensor) / fold(mask_t)
output_tensor = output_tensor[:, :, :output_tensor.shape[2] -
padding[0], :output_tensor.shape[3] -
padding[1]]
return output_tensor
def __init__(self,
device,
mask,
imgsz,
kernel_size,
stride,
output_size,
bias=True):
super(lp_pooling2d, self).__init__()
self.mask = Parameter(mask, requires_grad=True)
self.p_norm = Parameter(torch.zeros(output_size).add_(4),
requires_grad=True)
self.eps = 1e-60
self.sigmoid = nn.Sigmoid()
self.temperture = 5
self.pooling_operation = "Non_LSE"
self.unfold = nn.Unfold(kernel_size=(kernel_size, kernel_size),
stride=stride)
self.fold = nn.Fold(output_size=(imgsz // kernel_size,
imgsz // kernel_size),
kernel_size=(1, 1))
if bias:
self.bias = Parameter(torch.Tensor(output_size))
self.ondo_w = "True"
self.exp_p = "True"
def pat2im(patches, img_shape, patch_size):
"""
uniform aggregation
"""
y = patches.transpose(0, 1).unsqueeze(0)
uns = torch.ones(y.shape, device=y.device)
myfold = nn.Fold(output_size=img_shape, kernel_size=patch_size)
return myfold(y) / myfold(uns)
def compute_heatmap(self, img: torch.Tensor, patch_size, unfold_stride):
"""
img.shape -> (B, C, iH, iW)
B : batch size
C : channels of image (same to patches.shape[1])
iH : height of image
iW : width of image
patches.shape -> (P, C, pH, pW)
P : patch size
C : channels of image (same to img.shape[1])
pH : height of patch
pW : width of patch
"""
patches = self.patchize(img, patch_size, unfold_stride)
B, C, iH, iW = img.shape
P, C, pH, pW = patches.shape
heatmap = torch.zeros(P)
quotient, remainder = divmod(P, self.pbatch_size)
for i in range(quotient):
start = i * self.pbatch_size
end = start + self.pbatch_size
patch = patches[start:end, :, :, :]
patch = patch.to(self.device)
patch1, patch2 = torch.split(patch, [4, 4], dim=1)
surrogate_label = self.online_network(patch1)
pred = self.online_network(patch2)
losses = self.compute_squared_l2_distance(pred, surrogate_label)
heatmap[start:end] = losses
if remainder != 0:
patch = patches[-remainder:, :, :, :]
patch = patch.to(self.device)
patch1, patch2 = torch.split(patch, [4, 4], dim=1)
surrogate_label = self.online_network(patch1)
pred = self.online_network(patch2)
losses = self.compute_squared_l2_distance(pred, surrogate_label)
heatmap[-remainder:] = losses
fold = nn.Fold(
output_size=(iH, iW),
kernel_size=(pH, pW),
stride=unfold_stride,
)
heatmap = heatmap.expand(B, pH * pW, P)
heatmap = fold(heatmap)
heatmap = heatmap.squeeze()
del patches
return heatmap
def compute_heatmap(self,
img: T.Tensor) -> T.NDArray[(T.Any, T.Any), float]:
"""
img.shape -> (B, C, iH, iW)
B : batch size
C : channels of image (same to patches.shape[1])
iH : height of image
iW : width of image
patches.shape -> (P, C, pH, pW)
P : patch size
C : channels of image (same to img.shape[1])
pH : height of patch
pW : width of patch
"""
patches = self.patchize(img)
B, C, iH, iW = img.shape
P, C, pH, pW = patches.shape
heatmap = torch.zeros(P)
quotient, remainder = divmod(P, self.cfg.run.test.batch_size)
for i in range(quotient):
start = i * self.cfg.run.test.batch_size
end = start + self.cfg.run.test.batch_size
patch = patches[
start:
end, :, :, :] # (self.cfg.run.test.batch_size, C, pH, pW)
patch = patch.to(self.cfg.device)
surrogate_label, pred = self.school(patch)
losses = self.compute_squared_l2_distance(pred, surrogate_label)
heatmap[start:end] = losses
patch = patches[-remainder:, :, :, :]
patch = patch.to(self.cfg.device)
surrogate_label, pred = self.school(patch)
losses = self.compute_squared_l2_distance(pred, surrogate_label)
heatmap[-remainder:] = losses
fold = nn.Fold(
output_size=(iH, iW),
kernel_size=(pH, pW),
stride=self.cfg.run.test.unfold_stride,
)
heatmap = heatmap.expand(B, pH * pW, P)
heatmap = fold(heatmap)
heatmap = heatmap.squeeze().numpy()
del patches
return heatmap
def forward(self, x, sigma_map=None):
''' forward declaration '''
import numpy as np
sh = x.shape
# sigma may be missing
if sigma_map == None:
sigma_map = torch.ones(1, 1, sh[2], sh[3],
device=x.device) * self.sigma
# sigma may be a number
elif type(sigma_map) is not np.ndarray:
#print('scalar sigma map %f' % (np.array(sigma_map)*255))
sigma_map = torch.ones(1, 1, sh[2], sh[3],
device=x.device) * sigma_map
# subsample sigma_map
sig = sigma_map[:, :, ::2, ::2]
# unfold the image in 4 planes
#x = dtype(np.random.randn(1,1,4,4))
dx, dy = 0, 0
if sh[3] % 2 == 1: dx = 1
if sh[2] % 2 == 1: dy = 1
x = nn.functional.pad(x, (0, dx, 0, dy))
shtmp = x.shape
y = nn.Unfold(kernel_size=(2, 2), stride=2)(x)
#print(y.numpy().shape)
#print(x,y.reshape((1,4,2,2)))
y = y.reshape(sh[0], 4, shtmp[2] // 2, shtmp[3] // 2)
# channels are sorted differently in the pretrained FFDNET weights
indices = torch.tensor([0, 2, 1, 3], device=x.device)
y = torch.index_select(y, 1, indices)
# concatenate sigma_map with image
y = torch.cat([y, sig], dim=1)
# call dncnn
out = self.dncnn(y)
# fold the output
shout = out.shape
# channels are sorted differently in the pretrained FFDNET weights
#indices = torch.tensor([0, 2, 1, 3])
out = torch.index_select(out, 1, indices)
out = out.reshape(shout[0], 4, shout[2] * shout[3])
out = nn.Fold(output_size=(shout[2] * 2, shout[3] * 2),
kernel_size=(2, 2),
stride=2)(out)
return (out[:, :, :sh[2], :sh[3]])
def __init__(self, cvae_small, cvae_input_sz=16, stride=1, padding=0, same=True, img_size=64):
super(ConvVAE2d, self).__init__()
self.cvae_small = cvae_small
self.k = _pair(cvae_input_sz)
self.stride = _pair(stride)
self.padding = _quadruple(padding) # convert to l, r, t, b
self.same = same
self.fold = nn.Fold(output_size=(img_size, img_size), kernel_size=self.k, stride=self.stride)
self.padding = self._padding(size=(img_size, img_size))
def _get_patch_stiching(self, output_size: torch.Tensor):
"""stich non overlaping patches of size KERNEL_SIZE*KERNEL_SIZE.
Parameters:
output_size -tensor, dim-(1,2). represeting the |H|W| of the output
from all stiched patches
"""
return nn.Fold(
output_size=output_size,
kernel_size=self.KERNEL_SIZE,
stride=self.KERNEL_SIZE,
)
def forward(self, data):
fold = nn.Fold((1, data.shape[1]), (1, self.chunk_size),
stride=(1, self.hop_size))
unfold = nn.Unfold((1, self.chunk_size), stride=(1, self.hop_size))
batch_size, sequence_length, n_features, n_channels = (data.shape)
# extract chunks
data = data.transpose(1, 2)
data = self.bottleneck_norm(data)
data = data.transpose(1, -1)
data = self.bottleneck(data)
data = data.permute(0, 2, 3, 1)
data = data.reshape(batch_size, sequence_length, 1, -1)
data = data.transpose(3, 1)
data = unfold(data)
# unfold makes the data (batch_size, bottleneck_size * chunk_size, n_chunks)
n_chunks = data.shape[-1]
data = data.reshape(batch_size, -1, self.chunk_size, n_chunks)
data = data.transpose(3, 1)
# data is now (batch_size, n_chunks, chunk_size, bottleneck_size)
# process
output = data # Skip connection --->
for layer in self.layers: # |
data = layer(output) # |
if self.skip_connection: # |
output += data # <----<----<
else:
output = data
data = output
data = self.prelu(data)
# data is still (batch_size, n_chunks, chunk_size, bottleneck_size)
data = self.inv_bottleneck(data)
# data is now (batch_size, n_chunks, chunk_size, in_features)
data = data.transpose(1, -1)
data = self.output_norm(data)
data = data.transpose(1, -1)
# resynthesize with overlap/add
data = data.transpose(1, 3)
data = data.reshape(-1, n_features * self.chunk_size, n_chunks)
data = fold(data)
data = data.transpose(3, 1)
data = data.reshape(batch_size, sequence_length, n_features,
n_channels)
# renormalize after overlap/add
data = data / (self.chunk_size / self.hop_size)
return data
def inverse(self, data, **kwargs):
ndim = data.ndim
if ndim > 4:
# move sources to the batch dimension
# then fix it later
num_sources = data.shape[-1]
data = data.permute(0, -1, 1, 2, 3)
data = data.reshape(-1, *data.shape[2:])
data = self.apply_filter(data, **kwargs)
data *= self.window.sum()
data = data * self.window
num_batch, sequence_length, num_features, num_audio_channels = (
data.shape)
data = data.permute(0, 3, 2, 1)
data = data.reshape(-1, data.shape[2], data.shape[3])
fold = nn.Fold(
(1, self.output_length),
kernel_size=(1, self.filter_length),
stride=(1, self.hop_length),
dilation=self.dilation,
padding=(0, 0),
)
norm = data.new_ones(data.shape)
norm *= self.window.view(1, -1, 1)**2
data = fold(data)
norm = fold(norm)
norm[norm < 1e-10] = 1
data = data / norm
data = data.reshape(num_batch, num_audio_channels, -1)
boundary = self.filter_length // 2
data = data[..., boundary:-boundary]
data = data[..., :self.original_length]
if ndim > 4:
# then we moved sources to the batch dimension
# we need to move it back before returning
data = data.reshape(-1, num_sources, num_audio_channels,
data.shape[-1])
data = data.permute(0, 2, 3, 1)
return data
def backward(ctx, grad_output):
"""
Backward propagation of convolution operation
Args:
grad_output: gradients of the outputs
Outputs:
grad_input: gradients of the input features
grad_weight: gradients of the convolution weight
grad_bias: gradients of the bias term
"""
# unpack tensors and initialize the grads
input_unfolded, weight, bias = ctx.saved_tensors
grad_input = grad_weight = grad_bias = None
# recover the conv params
kernel_size = weight.size(2)
stride = ctx.stride
padding = ctx.padding
input_height = ctx.input_height
input_width = ctx.input_width
#################################################################################
# Fill in the code here
#################################################################################
# compute the gradients w.r.t. input and params
unfold_grad_op = nn.Unfold(kernel_size=(1, 1))
weight_unfolded = weight.view(weight.size(0), -1).transpose(0, 1)
grad_op_unfolded = unfold_grad_op(grad_output)
dx_unfolded = weight_unfolded.matmul(grad_op_unfolded)
fold_dx = nn.Fold(output_size=(input_height, input_width),
kernel_size=(kernel_size, kernel_size),
stride=stride,
padding=padding)
dx = fold_dx(dx_unfolded)
dw_unfolded = grad_op_unfolded.matmul(input_unfolded.transpose(1, 2))
dw = dw_unfolded.sum(dim=0)
dw_final = dw.view(weight.size())
if bias is not None and ctx.needs_input_grad[2]:
#compute the gradients w.r.t. bias (if any)
grad_bias = grad_output.sum((0, 2, 3))
return dx, dw_final, grad_bias, None, None
def __init__(self,
input_size,
num_channels,
num_filters,
batch_size,
kernel_size,
learning_rate,
f,
df,
inference_lr,
padding=0,
stride=1,
device="cpu",
numerical_test=False):
self.input_size = input_size
self.num_channels = num_channels
self.num_filters = num_filters
self.batch_size = batch_size
self.kernel_size = kernel_size
self.padding = padding
self.stride = stride
self.output_size = math.floor(
(self.input_size +
(2 * self.padding) - self.kernel_size) / self.stride) + 1
print(self.output_size)
self.learning_rate = learning_rate
self.inference_lr = inference_lr
self.f = f
self.df = df
self.device = device
self.weights = torch.empty(self.num_filters, self.num_channels,
self.kernel_size, self.kernel_size).normal_(
mean=0, std=0.05).to(self.device)
self.unfold = nn.Unfold(kernel_size=(self.kernel_size,
self.kernel_size),
padding=self.padding,
stride=self.stride).to(self.device)
self.fold = nn.Fold(output_size=(self.input_size, self.input_size),
kernel_size=(self.kernel_size, self.kernel_size),
padding=self.padding,
stride=self.stride).to(self.device)
self.numerical_test = numerical_test
self.use_conv_backwards_weights = False
self.use_conv_backwards_nonlinearity = False
self.update_backwards_weights = True
if self.use_conv_backwards_weights:
self.backwards_weights = torch.empty(
self.weights.reshape(self.num_filters, -1).T.shape).normal_(
mean=0, std=0.05).to(self.device)
def convmtx2_torch(H=torch.ones(3, 3), M=5, N=5):
P, Q = H.size()[0], H.size()[1]
rc = int((P - 1) / 2)
rw = int((Q - 1) / 2)
X = torch.ones(1, 1, M, N).cuda()
T = torch.zeros(M * N, M * N).cuda()
unfold = nn.Unfold(kernel_size=(P, Q), padding=(rc, rw))
output = unfold(X)
output_diag = torch.diag_embed(output, offset=0, dim1=-2, dim2=-1)
fold = nn.Fold(output_size=(M, N), kernel_size=(P, Q), padding=(rc, rw))
# input = torch.randn(1, 3 * 3 * 3, 1)
output2 = fold(
output_diag.permute(0, 3, 1, 2).contiguous().view(
(1, P * Q * M * N, M * N)))
def __init__(self,
in_channels,
out_channels,
kernel_size,
stride=(1, 1),
padding=(0, 0),
bias=True,
is_variational=True):
""" Constructor of the class, following the same syntax as the
PyTorch Conv2D class, with an optional is_variational
parameter.
Note that by default it is variational. You can switch on
and off this behavior by modifying the value of the attribute
is_variational.
"""
super(VarConv2d, self).__init__()
self.in_channels = in_channels
self.out_channels = out_channels
self.bias = bias
if isinstance(kernel_size, int):
self.kernel_size = (kernel_size, kernel_size)
else:
self.kernel_size = kernel_size
if isinstance(padding, int):
self.padding = (padding, padding)
else:
self.padding = padding
if isinstance(stride, int):
self.stride = (stride, stride)
else:
self.stride = stride
nb_in = self.kernel_size[0] * self.kernel_size[1] * self.in_channels
self.mu_layer = nn.Linear(nb_in, self.out_channels, bias)
self.logvar_layer = nn.Linear(nb_in, self.out_channels, bias)
self.unfold = nn.Unfold(kernel_size=self.kernel_size,
stride=self.stride,
padding=self.padding)
self.fold = nn.Fold(output_size=out_channels,
kernel_size=self.kernel_size,
stride=self.stride,
padding=self.padding)
self.is_variational = is_variational
def __init__(self,
input_size,
num_channels,
num_filters,
batch_size,
kernel_size,
learning_rate,
f,
df,
purely_positive=False,
use_layer_norm=True,
padding=0,
stride=1,
device="cpu"):
super(ConvLayer, self).__init__()
self.input_size = input_size
self.num_channels = num_channels
self.num_filters = num_filters
self.batch_size = batch_size
self.kernel_size = kernel_size
self.padding = padding
self.stride = stride
self.output_size = math.floor(
(self.input_size +
(2 * self.padding) - self.kernel_size) / self.stride) + 1
self.learning_rate = learning_rate
self.f = f
self.df = df
self.device = device
self.kernel = nn.Parameter(
torch.empty(self.num_filters, self.num_channels, self.kernel_size,
self.kernel_size).normal_(mean=0, std=0.05))
self.unfold = nn.Unfold(kernel_size=(self.kernel_size,
self.kernel_size),
padding=self.padding,
stride=self.stride)
self.fold = nn.Fold(output_size=(self.input_size, self.input_size),
kernel_size=(self.kernel_size, self.kernel_size),
padding=self.padding,
stride=self.stride)
self.purely_positive = purely_positive
self.use_layer_norm = use_layer_norm
self.layer_norm = nn.LayerNorm((self.batch_size, self.num_filters,
self.output_size, self.output_size))
self.bias = nn.Parameter(
torch.zeros((self.batch_size, self.num_filters, self.output_size,
self.output_size)))
def __init__(self, block_size, image_shape, overlapping=False):
"""
A module that extracts spatial patches from a 6D array with size [1, x, y, t, e, 2].
Output is also a 6D array with size [N, block_size, block_size, t, e, 2].
"""
super().__init__()
# Get image / block dimensions
self.block_size = block_size
self.image_shape = image_shape
_, self.nx, self.ny, self.nt, self.ne, _ = image_shape
# Overlapping vs. non-overlapping block settings
if overlapping:
block_stride = self.block_size // 2
# Use Hanning window to reduce blocking artifacts
win1d = torch.hann_window(block_size, dtype=torch.float32)**0.5
self.win = win1d[None, :, None, None, None,
None] * win1d[None, None, :, None, None, None]
else:
block_stride = self.block_size
self.win = torch.tensor([1.0], dtype=torch.float32)
# Figure out padsize (to avoid black bars)
num_blocks_x = (self.nx // self.block_size) + 2
num_blocks_y = (self.ny // self.block_size) + 2
self.pad_x = (self.block_size * num_blocks_x - self.nx) // 2
self.pad_y = (self.block_size * num_blocks_y - self.ny) // 2
nx_pad = self.nx + 2 * self.pad_x
ny_pad = self.ny + 2 * self.pad_y
# Compute total number of blocks
num_blocks_x = (self.nx - self.block_size +
2 * self.pad_x) / block_stride + 1
num_blocks_y = (self.ny - self.block_size +
2 * self.pad_y) / block_stride + 1
self.num_blocks = int(num_blocks_x * num_blocks_y)
# Set fold params
self.fold_params = dict(kernel_size=2 * (block_size, ),
stride=block_stride)
self.unfold_op = nn.Unfold(**self.fold_params)
self.fold_op = nn.Fold(output_size=(ny_pad, nx_pad),
**self.fold_params)
def one_kernel_feature_extraction(model: nn.Module,
input: torch.tensor,
D: int = 3,
padding: int = 1,
stride: int = 1) -> torch.tensor:
""" Extract features from intermediate layers
Return dim [F * N * input_dim]
Last frame is in accurate due to looping, proceed with caution
Args:
model (nn.Module): The model in question, Resnet 50
"""
# Transform input [F * N * C * H * W] -> [(F * N) * C * H * W]
F, N, C, H, W = input.size()
x = input.view(F * N, C, H, W)
unfolder = nn.Unfold(kernel_size=D, padding=padding, stride=stride)
# Compute correlation
correlations = []
# In order to get intermidiate features, we cannot use forward hook since this gives parallel issues
for i, part in enumerate(model.children()):
x = part(x)
# 2 features in total, each is [(F * N) * C_layer * H_layer * W_layer]
if i == 6:
# Shift output in F dimension, align each on with next frame
_, C_layer, H_layer, W_layer = x.size()
feature = x.view(F, N, C_layer, H_layer, W_layer)
next_feature = feature.roll(-1, dims=0)
# Each feature perform inner product with a block of D^2
feature = feature.unsqueeze(dim=3)
feature = feature.repeat(1, 1, 1, D**2, 1, 1).view(
(F * N, -1, H_layer * W_layer))
next_feature = unfolder(
next_feature.view(F * N, C_layer, H_layer, W_layer))
folder = nn.Fold(output_size=(H_layer, W_layer),
kernel_size=D,
padding=1)
# print(feature.size(), next_feature.size())
x = folder(feature * next_feature)
# Sum over C dimension, view into [F * N * (H_layer * W_layer)]
elif i > 7:
break
# Cat and return
return x.view(F, N, -1)
def __init__(self, args, mask, imgsz, kernel_size, stride, p, bias=True):
super(ShareConv2d, self).__init__()
self.args = args.kernel_size
self.mask = Parameter(mask, requires_grad=True)
self.p_norm = Parameter(torch.zeros(p).add_(4), requires_grad=True)
self.eps = 1e-60
self.sigmoid = nn.Sigmoid()
self.temperture = 5
self.pooling_operation = "LSE"
self.unfold = nn.Unfold(kernel_size=(kernel_size, kernel_size), stride=stride)
self.fold = nn.Fold(output_size=((args.imgsz // args.kernel_size), (args.imgsz // args.kernel_size)), kernel_size=(1, 1))
if bias:
self.bias = Parameter(torch.Tensor(p))
else:
self.register_parameter("bias", None)
self.ondo_w = "True"
self.exp_p = "True"
def __init__(self):
super(NNConvolutionModule, self).__init__()
self.input1d = torch.randn(1, 4, 36)
self.input2d = torch.randn(1, 4, 30, 10)
self.input3d = torch.randn(1, 4, 10, 4, 4)
self.module1d = nn.ModuleList([
nn.Conv1d(4, 33, 3),
nn.ConvTranspose1d(4, 33, 3),
nn.Fold(output_size=(5, 10), kernel_size=(2, 2)),
])
self.module2d = nn.ModuleList([
nn.Conv2d(4, 33, 3),
nn.ConvTranspose2d(4, 33, 3),
nn.Unfold(kernel_size=3),
])
self.module3d = nn.ModuleList([
nn.Conv3d(4, 33, 2),
nn.ConvTranspose3d(4, 33, 3),
])
def forward(self, x):
x_shape = x.shape
out_spatial_x = int(
math.floor((x_shape[3] -
(self.kern_size[0] - 1) - 1) / self.stride[0] + 1))
out_spatial_y = int(
math.floor((x_shape[4] -
(self.kern_size[1] - 1) - 1) / self.stride[1] + 1))
x = x.view(-1, self.in_channels, x_shape[3], x_shape[4])
temporal_buckets = nn.Unfold(kernel_size=self.kern_size,
stride=self.stride)(x).view(
x_shape[0], x_shape[1],
self.in_channels,
self.kern_size[0] * self.kern_size[1],
-1)
temporal_buckets_rot = temporal_buckets[:, 0, ...]
temporal_buckets_abs = temporal_buckets[:, 1, ...]
tbr_shape0 = temporal_buckets_rot.shape
temporal_buckets_rot = temporal_buckets_rot.permute(
0, 3, 1, 2).contiguous().view(-1, tbr_shape0[1], tbr_shape0[2])
temporal_buckets_abs = temporal_buckets_abs.permute(
0, 3, 1, 2).contiguous().view(-1, tbr_shape0[1], tbr_shape0[2])
tbr_shape = temporal_buckets_rot.shape
in_rot = temporal_buckets_rot * weightNormalize2(self.wmr)
in_abs = temporal_buckets_abs + weightNormalize1(self.wma)
in_rot = in_rot.view(tbr_shape0[0], out_spatial_x, out_spatial_y,
-1).permute(0, 3, 1, 2).contiguous().unsqueeze(1)
in_abs = in_abs.view(tbr_shape0[0], out_spatial_x, out_spatial_y,
-1).permute(0, 3, 1, 2).contiguous().unsqueeze(1)
in_ = torch.cat((in_rot, in_abs),
1).view(tbr_shape0[0], -1,
out_spatial_x * out_spatial_y)
in_fold = nn.Fold(output_size=(x_shape[3], x_shape[4]),
kernel_size=self.kern_size,
stride=self.stride)(in_)
in_fold = in_fold.view(x_shape[0], x_shape[1], x_shape[2], x_shape[3],
x_shape[4])
out = self.complex_conv(in_fold)
return out
def __init__(self, args):
super(FusionQualityMetric, self).__init__()
self.loss_type = 'Fusion'
self.WINDOW_SIZE = 8
self.out_size = tuple(
(torch.tensor(args.hr_shape) // self.WINDOW_SIZE))
self.unfold = nn.Unfold(kernel_size=(self.WINDOW_SIZE,
self.WINDOW_SIZE),
stride=self.WINDOW_SIZE)
self.fold = nn.Fold(output_size=self.out_size,
kernel_size=(1, 1),
stride=1)
self.convert1CH = transforms.Grayscale(num_output_channels=1)
self.toPIL = transforms.ToPILImage()
self.toTensor = transforms.ToTensor()
self.device = torch.device("cuda:0" if torch.cuda.is_available()
and args.device == "gpu" else "cpu")
def __init__(self, ch, resolution, kernel_size, mode="mlp"):
super(NeuralConvRecon, self).__init__()
self.kernel_size = kernel_size
self.linear1_out = (ch // kernel_size) * (kernel_size * kernel_size)
self.conv_ch_in = ch // kernel_size
self.interm_ch = self.linear1_out // (kernel_size * kernel_size)
self.mode = mode
# define architecture
if self.mode == "mlp":
self.linear1 = nn.Linear(ch, self.linear1_out // 2)
self.relu1 = nn.ReLU()
self.linear2 = nn.Linear(self.linear1_out // 2, self.linear1_out)
self.conv = nn.Conv2d(self.conv_ch_in,
ch,
kernel_size=1,
padding=0,
bias=False)
self.relu2 = nn.ReLU()
elif self.mode == "deconv":
# deconvolutionally fill in the holes
# upsample
# use use a weight to expand the area
#
pass
elif self.mode == "Intepolation":
# TODO: reconstruct by intepolate
pass
elif self.mode == "noconv":
# don't have convolution to avoid blur results?
pass
else:
raise NotImplementedError(
f"mode {self.mode} is not supported in NeuralConvRecon")
# fold
self.fold = nn.Fold(output_size=(resolution, resolution),
kernel_size=(kernel_size, kernel_size),
stride=3)
def __init__(self):
super().__init__()
self.conv = nn.Sequential(
nn.Conv2d(1, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.Conv2d(64, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
)
self.pool = nn.Sequential(
nn.MaxPool2d(kernel_size=pool_kernel,
stride=pool_kernel,
return_indices=True), )
self.unfold = nn.Sequential(
nn.Unfold(kernel_size=pool_kernel, stride=pool_kernel))
self.slidewindow = nn.Sequential(
nn.Unfold(kernel_size=smooth_kernel,
padding=smooth_padding,
stride=1))
self.fold = nn.Fold(output_size=(output_width, output_width),
kernel_size=(1, 1))
self.relu = nn.ReLU(inplace=True)
def forward(self, x, sigma_map=None):
''' forward declaration '''
sh = x.shape
# sigma may be missing
if sigma_map == None:
sigma_map = torch.ones(1, 1, sh[2], sh[3],
device=x.device) * self.sigma
# subsample sigma_map
sig = sigma_map[:, :, ::2, ::2]
# unfold the image in 4 planes
#x = dtype(np.random.randn(1,1,4,4))
dx, dy = 0, 0
if sh[3] % 2 == 1: dx = 1
if sh[2] % 2 == 1: dy = 1
x = nn.functional.pad(x, (0, dx, 0, dy))
shtmp = x.shape
y = nn.Unfold(kernel_size=(2, 2), stride=2)(x)
#print(y.numpy().shape)
#print(x,y.reshape((1,4,2,2)))
y = y.reshape(sh[0], 4, shtmp[2] // 2, shtmp[3] // 2)
# concatenate sigma_map with image
y = torch.cat([y, sig], dim=1)
# call dncnn
out = self.dncnn(y)
# fold the output
shout = out.shape
out = out.reshape(shout[0], 4, shout[2] * shout[3])
out = nn.Fold(output_size=(shout[2] * 2, shout[3] * 2),
kernel_size=(2, 2),
stride=2)(out)
return (out[:, :, :sh[2], :sh[3]])
def join(self, patches):
patches = patches.permute(1, 2, 3, 0)
prefolded = patches.reshape(1, -1, self.nx * self.ny)
folder = nn.Fold(
output_size = self.image_shape[2:],
kernel_size = self.kernel_size,
stride = self.stride
)
reconstructed = folder(prefolded)
# Divisor sorts out the normalisation; see torch.nn.Unfold documentation
if device == torch.device('cuda:0'):
im_ones = torch.cuda.FloatTensor(self.image_shape).fill_(1)
else:
im_ones = torch.ones(self.image_shape, dtype = float)
divisor = folder(self.unfolder(im_ones))
image = reconstructed / divisor
return self.cropImage(image)
def forward(self, input):
# input has shape (N, C, size1, size2)
# assert C == in_channel
N = input.size(0)
C = input.size(1)
s1 = input.size(2)
s2 = input.size(3)
Fold_f = nn.Fold((s1, s2), self.kernel_size, self.dilation,
self.padding, self.stride)
unfold = self.unf(input)
kernel_2 = self.kernel_size * self.kernel_size
L = unfold.size(2)
unfold = unfold.view(N, C, kernel_2, L)
unfold = unfold.permute(0, 3, 1, 2) # size (N, L, C, kernel_2)
after_max = torch.max(unfold,
self.thr).permute(0, 2, 3,
1).contiguous().view(N, -1, L)
fold_back = Fold_f(after_max)
return fold_back
def forward(ctx, input_feats, weight, bias, stride=1, padding=0):
"""
Forward propagation of convolution operation.
We only consider square filters with equal stride/padding in width and height!
Args:
input_feats: input feature map of size N * C_i * H * W
weight: filter weight of size C_o * C_i * K * K
bias: (optional) filter bias of size C_o
stride: (int, optional) stride for the convolution. Default: 1
padding: (int, optional) Zero-padding added to both sides of the input. Default: 0
Outputs:
output: responses of the convolution w*x+b
"""
# sanity check
assert weight.size(2) == weight.size(3)
assert input_feats.size(1) == weight.size(1)
assert isinstance(stride, int) and (stride > 0)
assert isinstance(padding, int) and (padding >= 0)
# save the conv params
kernel_size = weight.size(2)
ctx.stride = stride
ctx.padding = padding
ctx.input_height = input_feats.size(2)
ctx.input_width = input_feats.size(3)
# make sure this is a valid convolution
assert kernel_size <= (input_feats.size(2) + 2 * padding)
assert kernel_size <= (input_feats.size(3) + 2 * padding)
#################################################################################
# Fill in the code here
#################################################################################
# Calculating the output height and width
new_h = (input_feats.size(-2) + 2 * padding -
kernel_size) // stride + 1
new_w = (input_feats.size(-1) + 2 * padding -
kernel_size) // stride + 1
# Unfolding the inputs, weights, and output
unfold_input = nn.Unfold(kernel_size=kernel_size,
padding=padding,
stride=stride)
input_unfolded = unfold_input(input_feats)
weights_unfolded = weight.view(weight.size(0), -1)
output_unfolded = weights_unfolded.matmul(input_unfolded)
# Adding bias
op_unf = output_unfolded.transpose(0, 1)
for i in range(op_unf.size(0)):
op_unf[i] += bias[i]
output_unfolded = op_unf.transpose(0, 1)
fold_output = nn.Fold(output_size=(new_h, new_w), kernel_size=(1, 1))
output = fold_output(output_unfolded)
# save for backward (you need to save the unfolded tensor into ctx)
# ctx.save_for_backward(your_vars, weight, bias)
ctx.save_for_backward(input_unfolded, weight, bias)
return output
