def forward(self, x):
for t in range(4):
if t < 1:
Z_1 = self.W_b_1(x)
Z_2_mpool, indices_hid = self.MPool(Z_1)
Z_2 = self.W_b_2(self.LRN(F.relu(Z_2_mpool)))
read_out, indices_max = F.max_pool2d_with_indices(
self.LRN(F.relu(Z_2)),
kernel_size=Z_2.size()[2:],
return_indices=True)
final_z = self.Lin(read_out.view(-1, 32))
if t >= 1:
Z_1 = self.W_b_1(x) + self.W_l_1(self.LRN(
F.relu(Z_1))) + self.W_t_1(self.LRN(F.relu(Z_2)))
Z_2_mpool, indices_hid = self.MPool(Z_1)
Z_2 = self.W_b_2(self.LRN(F.relu(Z_2_mpool))) + self.W_l_2(
self.LRN(F.relu(Z_2)))
read_out, indices_max = F.max_pool2d_with_indices(
self.LRN(F.relu(Z_2)),
kernel_size=Z_2.size()[2:],
return_indices=True)
final_z = self.Lin(read_out.view(-1, 32))
#print(torch.sum(torch.isnan(final_z)))
#print(F.sigmoid(final_z[0,:]))
#print(final_z.size())
return (final_z)
def infer(self, data_dir, out_dir, shape, stride, device='cpu', epsilon=-1):
#Assume that the data_dir will contain pickled files which are 1 indexed
#Shape will be desired shape of filters
#Each file will be of shape [channels, x, y]
#Perform convolution then pooling
filters = self.f1.weight.data
#Prune weights
if epsilon != -1:
filters = prune_weights(filters, epsilon)
filters = torch.reshape(filters, shape).to(device)
files = [x for x in os.listdir(data_dir) if str.isdigit(x)]
for file in files:
X = torch.load(os.path.join(data_dir, file), map_location=device)
X_shape = X.shape
#Normalize input
X = s.normalize_single_sample(X.view(-1).unsqueeze(0))
X = torch.reshape(X, [1, X_shape[0], X_shape[1], X_shape[2]]).float()
#Convolve
out_map = F.conv2d(X, filters, stride=stride)
#Pooling
out_map, _ = F.max_pool2d_with_indices(out_map, 2, stride=1)
#Reshape to remove batch dimension
out_map = out_map.squeeze(0)
torch.save(out_map, os.path.join(out_dir, file))
def abs_maxpool_2d(x):
x_pool, index = F.max_pool2d_with_indices(torch.abs(x), (2, 2), 2)
x_flat = x.reshape([x.shape[0], x.shape[1], -1])
index_flat = index.reshape([index.shape[0], index.shape[1], -1])
x_selected = torch.gather(x_flat, 2, index_flat).reshape(x_pool.shape)
x_pool = x_pool * torch.sign(x_selected)
return x_pool
def compute_indices(self, x):
b, c, h, w = x.size()
oh, ow = 2 * h, 2 * w
x = -torch.arange(oh * ow, dtype=x.dtype, device=x.device).view(
1, 1, oh, ow)
indices = F.max_pool2d_with_indices(x, (2, 2),
(2, 2))[1].expand(b, c, h, w)
return indices
def forward(self, u, v):
B, C, _, _ = u.shape
x = vector2scalar(u, v)
_, idx = F.max_pool2d_with_indices(x,
kernel_size=self.kernel_size,
stride=self.stride,
padding=self.padding)
u = torch.gather(u.view(B, C, -1), 2, idx.view(B, C, -1)).view_as(idx)
v = torch.gather(v.view(B, C, -1), 2, idx.view(B, C, -1)).view_as(idx)
return u, v
def forward(self, input):
pooling_indices = []
x = input
for layer in self.encoder.children():
x = layer(x)
x, indices = F.max_pool2d_with_indices(x, kernel_size=2)
pooling_indices = [indices, *pooling_indices]
for layer in self.decoder.children():
indices, *pooling_indices = pooling_indices
x = F.max_unpool2d(x, indices, kernel_size=2)
x = layer(x)
return x
def find_local_max(self, box_prob):
B, C, H, W = box_prob.shape
max_prob, idx = F.max_pool2d_with_indices(box_prob,
3,
1,
1,
return_indices=True)
max_prob = max_prob[0, 0]
box_prob = box_prob[0, 0]
is_local_max = torch.nonzero(box_prob == max_prob)
y, x = is_local_max[:, 0], is_local_max[:, 1]
idx = torch.argsort(-box_prob[y, x])
k = self.scatter_topk
y = y[idx[:k]]
x = x[idx[:k]]
return y.cpu().numpy(), x.cpu().numpy(), box_prob[y, x]
def infer_from_spectro(self, data_dir, out_dir, shape, stride, device='cpu', epsilon=-1):
filters = self.f1.weight.data
#Prune weights
if epsilon != -1:
filters = prune_weights(filters, epsilon)
filters = torch.reshape(filters, [-1, 1, shape[2], shape[3]]).to(device)
from skimage.io import imread
#Extract folder names with assumption that all folder names are integers
dirs = [x for x in os.listdir(data_dir) if str.isdigit(x)]
for i in range(len(dirs)):
folder_name = dirs[i]
#Read in image
img = torch.load(os.path.join(data_dir, folder_name, 'spectro'))
if img.shape[0] > img.shape[1]:
continue
img = torch.as_tensor(img)
img_shape = img.shape
img = s.normalize_single_sample(img.view(-1).unsqueeze(0))
img = torch.reshape(img, [1, 1, img_shape[0], img_shape[1]]).float()
img = img.to(device)
#Convolve
out_map = F.conv2d(img, filters, stride=stride)
#Pooling
out_map, _ = F.max_pool2d_with_indices(out_map, 2, stride=1)
#Reshape to remove batch dimension
out_map = out_map.squeeze(0)
#Save
torch.save(out_map, os.path.join(out_dir, folder_name))
features: B * C * H * W
mask: B * H * W
'''
assert features.shape[2:4] == mask.shape[
1:3], 'shape mismatch between features and mask'
b, c, h, w = features.shape
mask = mask.unsqueeze(1).repeat(1, c, 1, 1) # B*H*W -> B*1*H*W
features = features * mask
pooled_features = F.avg_pool2d(features, [h, w])
scale = (h * w) / mask.sum(-1).sum(-1).clamp(min=1.0)
scale = scale.unsqueeze(-1).unsqueeze(-1)
return pooled_features * scale
if __name__ == '__main__':
b, c, h, w = 1, 3, 4, 4
features = torch.rand(b, c, h, w)
mask = torch.zeros(b, h, w)
mask[:, 1:3, 1:3] = 1
print('features', features)
# print ('mask', mask)
a, b = F.max_pool2d_with_indices(features, [h, w])
print(a)
print(b)
# max = mask_global_max_pooling_2d(features, mask)
# avg = mask_global_avg_pooling_2d(features, mask)
# print(max)
# print(avg)
def forward(self, input):
x11 = F.relu(self.bn11(self.conv11(input)))
x12 = F.relu(self.bn12(self.conv12(x11)))
x1p, id1 = F.max_pool2d_with_indices(x12,
kernel_size=2,
stride=2,
return_indices=True)
x21 = F.relu(self.bn21(self.conv21(x1p)))
x22 = F.relu(self.bn22(self.conv22(x21)))
x2p, id2 = F.max_pool2d_with_indices(x22,
kernel_size=2,
stride=2,
return_indices=True)
x31 = F.relu(self.bn31(self.conv31(x2p)))
x32 = F.relu(self.bn32(self.conv32(x31)))
x33 = F.relu(self.bn33(self.conv33(x32)))
x3p, id3 = F.max_pool2d_with_indices(x33,
kernel_size=2,
stride=2,
return_indices=True)
x41 = F.relu(self.bn41(self.conv41(x3p)))
x42 = F.relu(self.bn42(self.conv42(x41)))
x43 = F.relu(self.bn43(self.conv43(x42)))
x4p, id4 = F.max_pool2d_with_indices(x43,
kernel_size=2,
stride=2,
return_indices=True)
x51 = F.relu(self.bn51(self.conv51(x4p)))
x52 = F.relu(self.bn52(self.conv52(x51)))
x53 = F.relu(self.bn53(self.conv53(x52)))
x5p, id5 = F.max_pool2d(x53,
kernel_size=2,
stride=2,
return_indices=True)
print(x5p.size(), id5.size())
# unpooling - conv - bn - activation
# - conv - bn - activation
# - conv - bn - activation
# -
x5d = F.max_unpool2d(x5p, id5, kernel_size=2, stride=2)
x53d = F.relu(self.bn53d(self.conv53d(x5d)))
x52d = F.relu(self.bn52d(self.conv52d(x53d)))
x51d = F.relu(self.bn51d(self.conv51d(x52d)))
x4d = F.max_unpool2d(x51d, id4, kernel_size=2, stride=2)
x43d = F.relu(self.bn43d(self.conv43d(x4d)))
x42d = F.relu(self.bn42d(self.conv42d(x43d)))
x41d = F.relu(self.bn41d(self.conv41d(x42d)))
x3d = F.max_unpool2d(x41d, id3, kernel_size=2, stride=2)
x33d = F.relu(self.bn33d(self.conv33d(x3d)))
x32d = F.relu(self.bn32d(self.conv32d(x33d)))
x31d = F.relu(self.bn31d(self.conv31d(x32d)))
x2d = F.max_unpool2d(x31d, id2, kernel_size=2, stride=2)
x22d = F.relu(self.bn22d(self.conv22d(x2d)))
x21d = F.relu(self.bn21d(self.conv21d(x22d)))
x1d = F.max_unpool2d(x21d, id1, kernel_size=2, stride=2)
x12d = F.relu(self.bn12d(self.conv12d(x1d)))
x11d = self.conv11d(x12d)
return x11d
mask = torch.where(d >= bins[-1],
torch.tensor([1.]).to(device),
torch.tensor([0.]).to(device))
masked_s = torch.where(d >= bins[-1], g, torch.tensor([0.]).to(device))
semivariance1[:, :,
-1] = masked_s.sum(dim=2) / (mask.sum(dim=2) + 1e-7)
semivariance1[:, :, -1][mask.sum(dim=2) == 0] = 0
loss = ((semivariance1 - semivariance0)**2).mean()
return loss
def retrieve_elements_from_indices(self, tensor, indices):
flattened_tensor = tensor.flatten(start_dim=2)
output = flattened_tensor.gather(
dim=2, index=indices.flatten(start_dim=2)).view_as(indices)
return output
if __name__ == "__main__":
import torch.nn.functional as F
x1 = torch.randn(6, 64, 256, 256)
x0, ind = F.max_pool2d_with_indices(x1, 2)
kl = KrigingLoss()
loss = kl(x0, ind, x1)