def divide(a, b, limit=1, tol=1e-6):
a = tp.tensor(a)
b = tp.tensor(b)
a_s, b_s = a.shape ^ b.shape
a = a.view(a_s)
b = b.view(b_s)
shape = tp.Size(max(x, y) for x, y in zip(a_s, b_s))
return tp.where(b.abs() < tol, limit * tp.ones(shape), a / tp.where(b.abs() < tol, tol * tp.ones(shape), b))
def dBspline(i, U):
i = tp.tensor(i)
U = tp.tensor(U)
return (tp.where(
i == -1, -3 * (1 - U)**2 / 6,
tp.where(
i == 0, 3 * U**2 / 2 - 2 * U,
tp.where(i == 1, (-3 * U**2 + 2 * U + 1) / 2,
tp.where(i == 2, 3 * U**2 / 6, tp.zeros_like(U))))))
def Bspline(i, U):
i = tp.tensor(i)
U = tp.tensor(U)
return (tp.where(
i == -1, (1 - U)**3 / 6,
tp.where(
i == 0, U**3 / 2 - U * U + 2 / 3,
tp.where(i == 1, (-3 * U**3 + 3 * U * U + 3 * U + 1) / 6,
tp.where(i == 2, U**3 / 6, tp.zeros_like(U))))))
def __init__(self,
num_features: int,
eps: float = 1e-5,
momentum: float = 0.1,
affine: bool = True,
track_running_stats: bool = True) -> None:
super(_NormBase, self).__init__()
self.num_features = num_features
self.eps = eps
self.momentum = momentum
self.affine = affine
self.track_running_stats = track_running_stats
if self.affine:
self.weight = Parameter(torch.Tensor(num_features))
self.bias = Parameter(torch.Tensor(num_features))
else:
self.register_parameter('weight', None)
self.register_parameter('bias', None)
if self.track_running_stats:
self.register_buffer('running_mean', torch.zeros(num_features))
self.register_buffer('running_var', torch.ones(num_features))
self.register_buffer('num_batches_tracked',
torch.tensor(0, dtype=torch.long))
else:
self.register_parameter('running_mean', None)
self.register_parameter('running_var', None)
self.register_parameter('num_batches_tracked', None)
self.reset_parameters()
def to_image(data: Array,
nslice: [int, null] = None,
dim: int = -1,
has_cmap=False):
data = tp.tensor(data).squeeze()
if data.ndim <= 1:
raise TypeError(
"Please don't use 'plot.imshow' to demonstrate an array or a scalar. "
)
if data.nspace > 3:
raise TypeError(
f"'plot.imshow' takes 2 or 3D-data as input (currently {data.shape}), please reduce the dimension manually or specify special dimensions to reduce. "
)
if data.nspace == 3:
if data.has_batch: data = data.sample(random=False, dim=[])
if data.has_channel: data = data.sample(random=False, dim={})
if nslice is None:
if data.space[-1] <= 3: pass
elif data.space[0] <= 3: data = data.mvdim(0, 2)
else:
nslice = data.space[-1] // 2
data = data.pick(nslice, dim)
else:
data = data.pick(nslice, dim)
elif data.nspace == 2:
if data.has_batch: data = data.sample(random=False, dim=[])
if data.has_channel:
if has_cmap: data = data.sample(random=False, dim={})
else:
data = data.sample(number=min(data.channel_size, 3),
random=False,
dim={}).mvdim(data.channel_dimension, -1)
elif data.ndim == 3:
data = data.sample(random=False, dim=[])
return data.float().normalize()
def down_scale(image, *scaling:int):
image = tp.tensor(image)
if len(scaling) == 0:
scaling = (1,)
elif len(scaling) == 1 and iterable(scaling[0]):
scaling = scaling[0]
if len(scaling) == 1:
if isinstance(scaling[0], int):
scaling *= image.nspace
scaling = add_special(scaling, image.special, 1)
else: raise TypeError("Unknown scaling type for 'down_scale'. ")
elif len(scaling) < image.ndim and len(scaling) == image.nspace:
scaling = add_special(scaling, image.special, 1)
return image[tuple(slice(None, None, s) for s in scaling)]
def grad_image(array):
'''
Gradient image of array
array: (n_batch, n_feature, n_1, ..., n_{n_dim})
output: (n_batch, n_dim, n_feature, n_1, ..., n_{n_dim})
'''
array = tp.tensor(array)
output = tp.zeros_like(array)
grad_dim = int(array.has_batch)
output = []
for d in range(array.ndim):
if d in array.special: continue
b = (slice(None, None),) * d + (slice(2, None),) + (slice(None, None),) * (array.ndim - d - 1)
a = (slice(None, None),) * d + (slice(None, -2),) + (slice(None, None),) * (array.ndim - d - 1)
output.append(tp.crop_as((array[b] - array[a]) / 2, array))
return tp.stack(output, {grad_dim})
def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs):
version = local_metadata.get('version', None)
if (version is None or version < 2) and self.track_running_stats:
# at version 2: added num_batches_tracked buffer
# this should have a default value of 0
num_batches_tracked_key = prefix + 'num_batches_tracked'
if num_batches_tracked_key not in state_dict:
state_dict[num_batches_tracked_key] = torch.tensor(
0, dtype=torch.long)
super(_NormBase,
self)._load_from_state_dict(state_dict, prefix, local_metadata,
strict, missing_keys,
unexpected_keys, error_msgs)
def border(mask, min_length=10):
grid = tp.image_grid(*mask.shape)
mask = mask > 0.5
idx = mask[1:, :] ^ mask[:-1, :]
idx = idx.expand_to(2, -1, mask.size(1))
locs1 = (grid[:, 1:, :] + grid[:, :-1, :])[idx] / 2
idx = mask[:, 1:] ^ mask[:, :-1]
idx = idx.expand_to(2, mask.size(0), -1)
locs2 = (grid[:, :, 1:] + grid[:, :, :-1])[idx] / 2
locs = tp.cat(locs1.reshape(2, -1), locs2.reshape(2, -1), dim=1)
if locs.size == 0: return []
curves = []
unvisited = tp.ones(locs.shape[-1])
while True:
if not any(unvisited): break
first = tp.argmax(unvisited).item()
cloc = locs[:, first:first + 1]
unvisited[first] = 0
curve = cloc
while True:
dissq = tp.sum((locs - cloc)**2, 0)
inloc = tp.argmax(
tp.where((unvisited > 0) & (dissq > 0),
1 / dissq.clamp(min=0.1),
tp.tensor(0).float()))
if dissq[inloc] > 2: break
cloc = locs[:, inloc:inloc + 1]
curve = tp.cat(curve, cloc, dim=1)
unvisited[inloc] = 0
if not any(unvisited): break
sloc = locs[:, first:first + 1]
if tp.sum((cloc - sloc)**2) <= 2:
curve = tp.cat(curve, sloc, dim=1)
if curve.shape[1] <= min_length: continue
scurve = curve
for _ in range(100):
scurve = constraint(smooth(scurve), scurve, curve)
ccurve = scurve
for _ in range(100):
scurve = constraint(sharpen(scurve, curve), scurve, ccurve)
scurve = constraint(smooth(scurve), scurve, curve)
curves.append(scurve)
return curves
def clip_grad_norm_(parameters: _tensor_or_tensors,
max_norm: float,
norm_type: float = 2.0) -> torch.Tensor:
r"""Clips gradient norm of an iterable of parameters.
The norm is computed over all gradients together, as if they were
concatenated into a single vector. Gradients are modified in-place.
Arguments:
parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
single Tensor that will have gradients normalized
max_norm (float or int): max norm of the gradients
norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
infinity norm.
Returns:
Total norm of the parameters (viewed as a single vector).
"""
if isinstance(parameters, torch.Tensor):
parameters = [parameters]
parameters = [p for p in parameters if p.grad is not None]
max_norm = float(max_norm)
norm_type = float(norm_type)
if len(parameters) == 0:
return torch.tensor(0.)
device = parameters[0].grad.device
if norm_type == inf:
total_norm = max(p.grad.detach().abs().max().to(device)
for p in parameters)
else:
total_norm = torch.norm(
torch.stack([
torch.norm(p.grad.detach(), norm_type).to(device)
for p in parameters
]), norm_type)
clip_coef = max_norm / (total_norm + 1e-6)
if clip_coef < 1:
for p in parameters:
p.grad.detach().mul_(clip_coef.to(p.grad.device))
return total_norm
def forward(ctx, I1, I2, nbin=100):
with tp.no_grad():
if hasattr(ctx, 'JH'): del ctx.JH
nbin = tp.tensor(nbin)
data_pair = tp.stack(I1.flatten(1), I2.flatten(1), dim={1})
nbatch, nhist, ndata = data_pair.ishape
indices = []
values = []
ctx.window = (tp.image_grid(4, 4) - 1).flatten(1).transpose(0, 1)
for shift in ctx.window:
# [nbatch] x {nhist} x ndata
hist_pos = data_pair * nbin
index = tp.clamp(
tp.floor(hist_pos).long() + shift, 0, nbin - 1)
batch_idx = tp.arange(nbatch).expand_to([nbatch], {1}, ndata)
index = tp.cat(batch_idx, index, 1)
value = Bspline(shift.expand_to(data_pair),
tp.decimal(hist_pos)).prod(1)
indices.append(index)
values.append(value)
# n_batch x (1 + n_hist) x (n_data x 4 ** n_hist)
Mindices = tp.cat(indices, -1)
# n_batch x (n_data x 4 ** n_hist)
Mvalues = tp.cat(values, -1)
# (1 + n_hist) x (n_batch x n_data x 4 ** n_hist)
indices = Mindices.transpose(0, 1).flatten(1)
# (n_batch x n_data x 4 ** n_hist)
values = Mvalues.flatten(0)
if tp.Device == tp.DeviceCPU: creator = torch.sparse.FloatTensor
else: creator = torch.cuda.sparse.FloatTensor
collected = creator(indices, values,
(nbatch, nbin, nbin)).to_dense()
collected = tp.Tensor(collected, batch_dim=0)
ctx.nbin = nbin
ctx.Ishape = I1.shape
ctx.data_pair = data_pair
ctx.JH = collected / ndata
return ctx.JH
def up_scale(image, *scaling:int):
image = tp.tensor(image)
if len(scaling) == 0:
scaling = (1,)
elif len(scaling) == 1 and iterable(scaling[0]):
scaling = scaling[0]
if len(scaling) == 1:
if isinstance(scaling[0], int):
scaling *= image.nspace
scaling = add_special(scaling, image.special, 1)
else: raise TypeError("Unknown scaling type for 'up_scale'. ")
elif len(scaling) < image.ndim and len(scaling) == image.nspace:
scaling = add_special(scaling, image.special, 1)
for i, s in enumerate(scaling):
image = (
image
.transpose(i, -1)
.unsqueeze(-1)
.repeat((1,) * image.ndim + (int(s),))
.flatten(-2)
.transpose(i, -1)
)
return image
torch.manual_seed(0)
with scope("test torch, gpu"):
t_ = torch.randn(3000, 400).to(tp.Device).requires_grad_(True)
a_ = t_
LP_ = torch.nn.Linear(400, 400).to(tp.Device)
for _ in range(10): a_ = LP_(a_).relu()
a_.sum().backward()
assert a.is_cuda is True
assert t.allclose(t_)
assert isinstance(t, tp.Tensor)
assert isinstance(a, tp.Tensor)
assert isinstance(LP.weight, tp.nn.Parameter)
assert isinstance(LP.bias, tp.nn.Parameter)
assert isinstance(tp.tensor(np.array([1., 2.])), tp.Tensor)
if torch.cuda.is_available():
assert a.is_cuda
assert t.is_cuda
assert tp.tensor(np.array([1., 2.])).is_cuda
tp.set_autodevice(False)
tp.manual_seed(0)
with scope("test tp, cpu"):
t = tp.randn(3000, 400, requires_grad=True)
a = t
LP = tp.nn.Linear(400, 400)
for _ in range(10): a = LP(a).relu()
a.sum().backward()
torch.manual_seed(0)
def maskshow(*masks,
on=None,
alpha=0.5,
nslice=None,
dim=-1,
stretch=False,
**kwargs):
global canvas
if on is not None:
if isinstance(on, (int, tuple)): background(*on)
elif isarray(on): canvas = to_image(on, nslice, dim)
elif isinstance(on, list): canvas = to_image(Tensor(on), nslice, dim)
else: raise TypeError("Unrecognized argument 'on' for 'maskshow'. ")
elif canvas is None:
canvas = (1., ) * 3
if len(masks) == 0: return imshow
alpha = totuple(alpha, len(masks))
new_masks = []
new_alpha = []
for m, a in zip(masks, alpha):
img = to_image(m, nslice, dim)
if img.ndim == 3:
new_masks.extend(x.squeeze(-1) for x in img.split(1, dim=dim))
new_alpha.extend([a] * img.size(dim))
else:
new_masks.append(img)
new_alpha.append(a)
color_mask_map = [
(to_RGB(c), m, a)
for c, m, a in zip(colors * (len(new_masks) // len(colors) +
1), new_masks, new_alpha)
]
color_mask_map.extend((to_RGB(c), m, alpha[0]) for c, m in kwargs.items())
if not stretch:
shapes = [m.ishape for _, m, _ in color_mask_map]
target_shape = shapes[0]
if len(set(shapes)) > 1 or not isinstance(
canvas, tuple) and target_shape != canvas.shape:
raise TypeError(
"Please use masks of the same size as the background image, "
"or use 'stretch=True' in 'maskshow' to automatically adjust the image sizes. "
)
else:
def adjust(m, to):
ms = tuple(m.shape)
scaling = tuple((a // b, b // a) for a, b in zip(to, ms))
return m.down_scale([max(v, 1) for u, v in scaling
]).up_scale([max(u, 1)
for u, v in scaling]).crop_as(to)
shapes = [m.ishape for _, m, _ in color_mask_map]
if not isinstance(canvas, tuple): shapes.append(canvas.shape[:2])
areas = [u * v for u, v in shapes]
target_shape = shapes[areas.index(max(areas))]
color_mask_map = [(c, adjust(m, to=target_shape), a)
for c, m, a in color_mask_map]
canvas = adjust(canvas, to=target_shape)
target_shape = tp.Size(*target_shape, {3})
if isinstance(canvas, tuple):
canvas = tp.tensor(list(canvas)).expand_to(target_shape)
elif canvas.ndim == 2:
canvas = canvas.expand_to(target_shape)
coeff = vector(1 - a * m for _, m, a in color_mask_map).prod()
canvas *= coeff
for i, (c, m, a) in enumerate(color_mask_map):
coeff = vector(a * m if j == i else 1 - a * m
for j, (_, m, a) in enumerate(color_mask_map)).prod()
canvas += coeff.unsqueeze(-1) * m.unsqueeze(-1) * tp.tensor(
list(c)).unsqueeze(0, 1)
return plt.imshow(canvas.numpy(), **kwargs)
