def oneMix(mask, data=None, target=None): #Mix if not (data is None): stackedMask0, _ = torch.broadcast_tensors(mask[0], data[0]) data = (stackedMask0 * data[0] + (1 - stackedMask0) * data[1]).unsqueeze(0) if not (target is None): stackedMask0, _ = torch.broadcast_tensors(mask[0], target[0]) target = (stackedMask0 * target[0] + (1 - stackedMask0) * target[1]).unsqueeze(0) return data, target
def face_add_arbitrary(faces, face_ix_to, per_face_data, output, counts=None): ind_1d = faces[:, face_ix_to] one = torch.ones((1, ), dtype=torch.float32, device=faces.device) if len(ind_1d.shape) < len(per_face_data.shape): ind_2d, _ = torch.broadcast_tensors(ind_1d.unsqueeze(1), per_face_data) else: ind_2d = ind_1d one_1d, _ = torch.broadcast_tensors(one, ind_1d) output.scatter_add_(0, ind_2d, per_face_data) if counts is not None: counts.scatter_add_(0, ind_1d, one_1d)
def __and__(self, y): """Bitwise AND operator (element-wise)""" result = self.clone() # TODO: Remove explicit broadcasts to allow smaller beaver triples if isinstance(y, BinarySharedTensor): broadcast_tensors = torch.broadcast_tensors(result.share, y.share) result.share = broadcast_tensors[0].clone() elif torch.is_tensor(y): broadcast_tensors = torch.broadcast_tensors(result.share, y) result.share = broadcast_tensors[0].clone() return result.__iand__(y)
def normalize(MEAN, STD, data=None, target=None): #Normalize if not (data is None): if data.shape[1] == 3: STD = torch.Tensor(STD).unsqueeze(0).unsqueeze(2).unsqueeze( 3).cuda() MEAN = torch.Tensor(MEAN).unsqueeze(0).unsqueeze(2).unsqueeze( 3).cuda() STD, data = torch.broadcast_tensors(STD, data) MEAN, data = torch.broadcast_tensors(MEAN, data) data = ((data - MEAN) / STD).float() return data, target
def test_method_interface(): shape = 11, 13 lin = LinearMasked(*shape) mask = torch.tensor(1.0) lin.mask_(mask) assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask)) assert torch.allclose(lin.weight_masked, mask * lin.weight) lin.mask_(None) assert lin.mask is None with pytest.raises(RuntimeError, match=r"has no sparsity mask"): lin.weight_masked # enable/disable mask = torch.randint(2, size=(1, 1)).float() lin.mask_(mask) assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask)) assert torch.allclose(lin.weight_masked, mask * lin.weight) # output masking mask = torch.randint(2, size=(shape[1], 1)).float() lin.mask_(mask) assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask)) assert torch.allclose(lin.weight_masked, mask * lin.weight) # input masking mask = torch.randint(2, size=( 1, shape[0], )).float() lin.mask_(mask) assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask)) assert torch.allclose(lin.weight_masked, mask * lin.weight) # unstructured masking mask = torch.randint(2, size=( shape[1], shape[0], )).float() lin.mask_(mask) assert torch.allclose(*torch.broadcast_tensors(lin.mask, mask)) assert torch.allclose(lin.weight_masked, mask * lin.weight) # mask is overwritten (recreated), so `torch.expand` raises with pytest.raises(RuntimeError, match=r"must match the existing size"): lin.mask_(torch.ones(shape[1], 2)) # mask is set anew, so `torch.expand` raises lin.mask_(None) with pytest.raises(RuntimeError, match=r"must match the existing size"): lin.mask_(torch.ones(0, shape[0]))
def forward(self, x): feature = self.layer(x).view(-1, self.numkernel, self.dimsize) feature = feature.unsqueeze(3) feature_ = feature.permute(3, 1, 2, 0) feature, feature_ = torch.broadcast_tensors(feature, feature_) norm = torch.sum(torch.abs(feature - feature_), dim=2) eraser = torch.eye(feature.shape[0], device=torch.device('cuda')).view( feature.shape[0], 1, feature.shape[0]) eraser, norm = torch.broadcast_tensors(eraser, norm) c_b = torch.exp(-(norm + eraser * 1e6)) o_b = torch.sum(c_b, dim=2) x = torch.cat((x, o_b), dim=1) return x
def tridiagonal_solve(b, A_upper, A_diagonal, A_lower): """Solves a tridiagonal system Ax = b. The arguments A_upper, A_digonal, A_lower correspond to the three diagonals of A. Letting U = A_upper, D=A_digonal and L = A_lower, and assuming for simplicity that there are no batch dimensions, then the matrix A is assumed to be of size (k, k), with entries: D[0] U[0] L[0] D[1] U[1] L[1] D[2] U[2] 0 L[2] D[3] U[3] . . . . . . . . . L[k - 3] D[k - 2] U[k - 2] 0 L[k - 2] D[k - 1] U[k - 1] L[k - 1] D[k] Arguments: b: A tensor of shape (..., k), where '...' is zero or more batch dimensions A_upper: A tensor of shape (..., k - 1). A_diagonal: A tensor of shape (..., k). A_lower: A tensor of shape (..., k - 1). Returns: A tensor of shape (..., k), corresponding to the x solving Ax = b Warning: This implementation isn't super fast. You probably want to cache the result, if possible. """ # This implementation is very much written for clarity rather than speed. A_upper, _ = torch.broadcast_tensors(A_upper, b[..., :-1]) A_lower, _ = torch.broadcast_tensors(A_lower, b[..., :-1]) A_diagonal, b = torch.broadcast_tensors(A_diagonal, b) channels = b.size(-1) new_b = np.empty(channels, dtype=object) new_A_diagonal = np.empty(channels, dtype=object) outs = np.empty(channels, dtype=object) new_b[0] = b[..., 0] new_A_diagonal[0] = A_diagonal[..., 0] for i in range(1, channels): w = A_lower[..., i - 1] / new_A_diagonal[i - 1] new_A_diagonal[i] = A_diagonal[..., i] - w * A_upper[..., i - 1] new_b[i] = b[..., i] - w * new_b[i - 1] outs[channels - 1] = new_b[channels - 1] / new_A_diagonal[channels - 1] for i in range(channels - 2, -1, -1): outs[i] = (new_b[i] - A_upper[..., i] * outs[i + 1]) / new_A_diagonal[i] return torch.stack(outs.tolist(), dim=-1)
def __init__(self, loc, covariance_matrix=None, precision_matrix=None, scale_tril=None, validate_args=None): if loc.dim() < 1: raise ValueError("loc must be at least one-dimensional.") if (covariance_matrix is not None) + (scale_tril is not None) + ( precision_matrix is not None) != 1: raise ValueError( "Exactly one of covariance_matrix or precision_matrix or scale_tril may be specified." ) loc_ = loc.unsqueeze(-1) # temporarily add dim on right if scale_tril is not None: if scale_tril.dim() < 2: raise ValueError( "scale_tril matrix must be at least two-dimensional, " "with optional leading batch dimensions") self.scale_tril, loc_ = torch.broadcast_tensors(scale_tril, loc_) elif covariance_matrix is not None: if covariance_matrix.dim() < 2: raise ValueError( "covariance_matrix must be at least two-dimensional, " "with optional leading batch dimensions") self.covariance_matrix, loc_ = torch.broadcast_tensors( covariance_matrix, loc_) else: if precision_matrix.dim() < 2: raise ValueError( "precision_matrix must be at least two-dimensional, " "with optional leading batch dimensions") self.precision_matrix, loc_ = torch.broadcast_tensors( precision_matrix, loc_) self.loc = loc_[..., 0] # drop rightmost dim self.normalizing_constant = torch.nn.Parameter(torch.tensor([1.])) batch_shape, event_shape = self.loc.shape[:-1], self.loc.shape[-1:] super(UnnormMVGaussian, self).__init__(batch_shape, event_shape, validate_args=validate_args) if scale_tril is not None: self._unbroadcasted_scale_tril = scale_tril else: if precision_matrix is not None: self.covariance_matrix = torch.inverse( precision_matrix).expand_as(loc_) self._unbroadcasted_scale_tril = torch.cholesky( self.covariance_matrix)
def __init__(self, loc, covariance_matrix=None, precision_matrix=None, scale_tril=None, validate_args=None): if loc.dim() < 1: raise ValueError("loc must be at least one-dimensional.") if (covariance_matrix is not None) + (scale_tril is not None) + ( precision_matrix is not None) != 1: raise ValueError( "Exactly one of covariance_matrix or precision_matrix or scale_tril may be specified." ) loc_ = loc.unsqueeze(-1) # temporarily add dim on right if scale_tril is not None: if scale_tril.dim() < 2: raise ValueError( "scale_tril matrix must be at least two-dimensional, " "with optional leading batch dimensions") self.scale_tril, loc_ = torch.broadcast_tensors(scale_tril, loc_) elif covariance_matrix is not None: if covariance_matrix.dim() < 2: raise ValueError( "covariance_matrix must be at least two-dimensional, " "with optional leading batch dimensions") self.covariance_matrix, loc_ = torch.broadcast_tensors( covariance_matrix, loc_) else: if precision_matrix.dim() < 2: raise ValueError( "precision_matrix must be at least two-dimensional, " "with optional leading batch dimensions") self.precision_matrix, loc_ = torch.broadcast_tensors( precision_matrix, loc_) self.loc = loc_[..., 0] # drop rightmost dim batch_shape, event_shape = self.loc.shape[:-1], self.loc.shape[-1:] super(MultivariateNormal, self).__init__(batch_shape, event_shape, validate_args=validate_args) if scale_tril is not None: self._unbroadcasted_scale_tril = scale_tril elif covariance_matrix is not None: #self._unbroadcasted_scale_tril = torch.cholesky(covariance_matrix) self._unbroadcasted_scale_tril = torch.sqrt(covariance_matrix) else: # precision_matrix is not None raise NotImplementedError( 'Only covariance_matrix or scale_tril may be specified')
def generate_binary_triple(size0, size1, device=None): """Generate binary triples of given size""" generator = TTPClient.get().get_generator(device=device) a = generate_kbit_random_tensor(size0, generator=generator, device=device) b = generate_kbit_random_tensor(size1, generator=generator, device=device) if comm.get().get_rank() == 0: # Request c from TTP c = TTPClient.get().ttp_request("binary", device, size0, size1) else: size2 = torch.broadcast_tensors(a, b)[0].size() c = generate_kbit_random_tensor(size2, generator=generator, device=device) # Stack to vectorize scatter function a = BinarySharedTensor.from_shares(a) b = BinarySharedTensor.from_shares(b) c = BinarySharedTensor.from_shares(c) return a, b, c
def log_prob(self, value): if self._validate_args: self._validate_sample(value) value = value.long().unsqueeze(-1) value, log_pmf = torch.broadcast_tensors(value, self.logits) value = value[..., :1] return log_pmf.gather(-1, value).squeeze(-1)
def forward(self, inputs, caps, mask=None): B, R, input_dim = inputs.shape # (B, R, in) C, output_dim = caps.shape[1], self.output_dim # inputs = inputs[None, :, :, None, :] # (1, B, R, 1, in) caps = caps[None, :, :, None, :].transpose(0, 2) # (C, B, 1, 1, in) inputs, caps = torch.broadcast_tensors(inputs, caps) # (C, B, R, 1, in) u = torch.cat([inputs, caps], axis=-1) # (C, B, R, 1, 2*in) u = self.dense_out(u) # (C, B, R, 1, out) if mask is not None: float_mask = mask.float()[None, :, :, None, None] # (B, R) => (1, B, R, 1, 1) else: float_mask = 1. b = torch.zeros(C, B, R, 1, 1).to(inputs.device) # (C, B, R, 1, 1) for i in range(self.n_iter): v = b.softmax(dim=0) * float_mask # (C, B, R, 1, 1) c_hat = (u * v).sum(2, keepdims=True) # (C, B, 1, 1, out) c = squash(c_hat, dim=-1) # (C, B, 1, 1, out) b = b + (c * u).sum(-1, keepdims=True) # (C, B, R, 1, 1) out_c = c[:, :, 0, 0, :].permute(1, 0, 2) # (B, C, out) out_v = v[:, :, :, 0, 0].permute(1, 2, 0) # (B, R, C) out_b = b[:, :, :, 0, 0].permute(1, 2, 0) # (B, R, C) # print(out_v[0, 0]) return out_c, out_v, out_b
def proj_tangent(x, u): assert x.shape[-2:] == u.shape[-2:], "Wrong shapes" x, u = torch.broadcast_tensors(x, u) x_shape = x.shape x = x.reshape(-1, x_shape[-2], x_shape[-1]) u = u.reshape(-1, x_shape[-2], x_shape[-1]) xt = x.transpose(-1, -2) batch_size, n = x.shape[0:2] I = torch.eye(n, dtype=x.dtype, device=x.device) I = I.expand_as(x) mu = x * u A = linalg.block_matrix([[I, x], [xt, I]]) B = A[:, :, 1:] z1 = mu.sum(dim=-1).unsqueeze(-1) zt1 = mu.sum(dim=-2).unsqueeze(-1) b = torch.cat( [z1, zt1], dim=1, ) rhs = B.transpose(1, 2) @ (b - A[:, :, 0:1]) lhs = B.transpose(1, 2) @ B zeta, _ = torch.solve(rhs, lhs) alpha = torch.cat( [torch.ones(batch_size, 1, 1, dtype=x.dtype), zeta[:, 0:n - 1]], dim=1) beta = zeta[:, n - 1:2 * n - 1] rgrad = mu - (alpha + beta.transpose(-1, -2)) * x rgrad = rgrad.reshape(x_shape) return rgrad
def __init__(self, loc, cov_factor, cov_diag, validate_args=None): if loc.dim() < 1: raise ValueError("loc must be at least one-dimensional.") event_shape = loc.shape[-1:] if cov_factor.dim() < 2: raise ValueError("cov_factor must be at least two-dimensional, " "with optional leading batch dimensions") if cov_factor.shape[-2:-1] != event_shape: raise ValueError("cov_factor must be a batch of matrices with shape {} x m" .format(event_shape[0])) if cov_diag.shape[-1:] != event_shape: raise ValueError("cov_diag must be a batch of vectors with shape {}".format(event_shape)) loc_ = loc.unsqueeze(-1) cov_diag_ = cov_diag.unsqueeze(-1) try: loc_, self.cov_factor, cov_diag_ = torch.broadcast_tensors(loc_, cov_factor, cov_diag_) except RuntimeError: raise ValueError("Incompatible batch shapes: loc {}, cov_factor {}, cov_diag {}" .format(loc.shape, cov_factor.shape, cov_diag.shape)) self.loc = loc_[..., 0] self.cov_diag = cov_diag_[..., 0] batch_shape = self.loc.shape[:-1] self._unbroadcasted_cov_factor = cov_factor self._unbroadcasted_cov_diag = cov_diag self._capacitance_tril = _batch_capacitance_tril(cov_factor, cov_diag) super(LowRankMultivariateNormal, self).__init__(batch_shape, event_shape, validate_args=validate_args)
def forward(self, x): set_feat = self.small_xf(x) comb_feat = torch.cat(torch.broadcast_tensors(x, set_feat.unsqueeze(-2)), dim=-1) heads = self.w(comb_feat).view(comb_feat.size()[:3]+(self.heads, -1)) q = self.q.view((1, 1, 1)+self.q.size()) e_heads = (heads * q).sum(dim=-1) return F.softmax(e_heads, dim=-2)
def _broadcast_prediction_target(prediction, target): """ broadcast prediction and target in identical shapes Args: prediction (Tensor): prediction target (Tensor): target Returns: prediction and target in the same shape """ if not torch.is_tensor(prediction): prediction = torch.tensor(prediction) if not torch.is_tensor(target): target = torch.tensor(target, dtype=prediction.dtype, device=prediction.device) target = target.to(dtype=prediction.dtype, device=prediction.device) while len(prediction.shape) < 4: prediction = prediction[:, None] while len(target.shape) < 4: target = target[:, None] try: prediction, target = torch.broadcast_tensors(prediction.clone(), target.clone()) except RuntimeError: raise RuntimeError("failed to broadcast target in the same shape as prediction") return prediction, target
def broadcast_all(*values): r""" Given a list of values (possibly containing numbers), returns a list where each value is broadcasted based on the following rules: - `torch.*Tensor` instances are broadcasted as per :ref:`_broadcasting-semantics`. - numbers.Number instances (scalars) are upcast to tensors having the same size and type as the first tensor passed to `values`. If all the values are scalars, then they are upcasted to scalar Tensors. Args: values (list of `numbers.Number` or `torch.*Tensor`) Raises: ValueError: if any of the values is not a `numbers.Number` or `torch.*Tensor` instance """ if not all( isinstance(v, torch.Tensor) or isinstance(v, Number) for v in values): raise ValueError( 'Input arguments must all be instances of numbers.Number or torch.tensor.' ) if not all([isinstance(v, torch.Tensor) for v in values]): options = dict(dtype=torch.get_default_dtype()) for value in values: if isinstance(value, torch.Tensor): options = dict(dtype=value.dtype, device=value.device) break values = [ v if isinstance(v, torch.Tensor) else torch.tensor(v, **options) for v in values ] return torch.broadcast_tensors(*values)
def color_loss(input, target, reduction='mean'): pixDictionary = buildPixDictionary(target) print("Len pixDict", len(pixDictionary)) print(pixDictionary) print(input.size()) batch, channel, height, width = input.size() ret = torch.empty(batch, channel, height, width) nbPixelDif = len(pixDictionary) nbPixel = height * width * batch if target.requires_grad: for img in range(batch): for i in range(height): for j in range(width): mse = (input[img, :, i, j] - target[img, :, i, j])**2 nbApparition = pixDictionary[getIndexPix( input[img, :, i, j].detach())] color_coef = nbPixel / (nbPixelDif * nbApparition) ret[img, :, i, j] = mse * color_coef else: expanded_input, expanded_target = torch.broadcast_tensors( input, target) for i in range(height): for j in range(width): mse = (expanded_input[i][j] - expanded_target[i][j])**2 nbApparition = pixDictionary[input[i][j]] color_coef = nbPixel / (nbPixelDif * nbApparition) ret[i][j] = mse * color_coef if reduction != 'none': ret = torch.mean(ret) if reduction == 'mean' else torch.sum(ret)
def __init__(self, name: str, shape: tc.Size, device=None, **params): # for distributions whose parameter and random variable (or, one sample) have the same shape if device is None: device = Distr.default_device fnnames, fnvals, tennames, tenvals = [], [], [], [] for pmname, pmval in params.items(): if callable(pmval): fnnames.append(pmname) fnvals.append(pmval) else: tennames.append(pmname) tenvals.append(pmval) tenvals = tensorify(device, *tenvals) if shape is None: if Distr.has_name(name): shape = Distr.shape_var(name) elif tenvals: shape = tc.broadcast_tensors(*tenvals)[0].shape else: shape = tc.Size() parents = set() for fname, fval in zip(fnnames, fnvals): parents_inc = fargnames(fval) if parents_inc: parents |= parents_inc setattr(self, fname, fedic(fval)) else: setattr(self, fname, lambda conds, fval=fval: tensorify(device, fval())[0]. expand(Distr.shape_bat(conds) + shape)) for tname, tval in zip(tennames, tenvals): setattr(self, tname, lambda conds, tval=tval: tval.expand( Distr.shape_bat(conds) + shape)) super(DistrElem, self).__init__(names_shapes={name: shape}, parents=parents) self._name, self._shape, self._device = name, shape, device
def broadcast_all(*values): r""" Given a list of values (possibly containing numbers), returns a list where each value is broadcasted based on the following rules: - `torch.*Tensor` instances are broadcasted as per :ref:`_broadcasting-semantics`. - numbers.Number instances (scalars) are upcast to tensors having the same size and type as the first tensor passed to `values`. If all the values are scalars, then they are upcasted to scalar Tensors. Args: values (list of `numbers.Number` or `torch.*Tensor`) Raises: ValueError: if any of the values is not a `numbers.Number` or `torch.*Tensor` instance """ if not all(torch.is_tensor(v) or isinstance(v, Number) for v in values): raise ValueError( 'Input arguments must all be instances of numbers.Number or torch.tensor.' ) if not all(map(torch.is_tensor, values)): new_tensor = _default_promotion for value in values: if torch.is_tensor(value): new_tensor = value.new_tensor break values = [v if torch.is_tensor(v) else new_tensor(v) for v in values] return torch.broadcast_tensors(*values)
def __init__(self, loc, scale=None, link="probit", sampling=5000, validate_args=None): if loc.dim() < 1: raise ValueError("loc must be at least one-dimensional.") loc_ = loc.unsqueeze(-1) # temporarily add dim on right if scale.dim() < 2: raise ValueError("scale must be at least two-dimensional, " "with optional leading batch dimensions") self.scale, loc_ = torch.broadcast_tensors(scale, loc_) self.scale = scale self._covariance_matrix = scale self.loc = loc_[..., 0] # drop rightmost dim self.df = scale.shape[-1] self.sampling = sampling self.alpha = 1.0 if link == "logit": self.link = lambda value: torch.sigmoid(value) self.alpha = 1.7012 elif link == "probit": self.link = lambda value: torch.distributions.Normal(0.0, 1.0).cdf( value) elif link == "linear": self.link = lambda value: torch.clamp(value, 0.0, 1.0) batch_shape, event_shape = self.loc.shape[:-1], self.loc.shape[-1:] super(MultivariateProbit, self).__init__(batch_shape, event_shape, validate_args=validate_args)
def pairwise_soft_dtw(data1, data2, sdtw=None, device=torch.device('cpu')): if sdtw is None: raise ValueError('sdtw is None - initialize it with SoftDTW') # transfer to device data1, data2 = data1.to(device), data2.to(device) # (batch_size, seq_len, feature_dim=1) A = data1.unsqueeze(dim=2) # (cluster_size, seq_len, feature_dim=1) B = data2.unsqueeze(dim=2) distances = [] for b in B: # (1, seq_len, 1) b = b.unsqueeze(dim=0) A, b = torch.broadcast_tensors(A, b) # (batch_size, 1) sdtw_distance = sdtw(b, A).view(-1, 1) distances.append(sdtw_distance) # (batch_size, cluster_size) dis = torch.cat(distances, dim=1) return dis
def forward(self, predicted_mesh): predicted_mesh_verts = predicted_mesh.vertices predicted_mesh_indices = predicted_mesh.indices lap = laplacian(predicted_mesh_verts, predicted_mesh_indices) # just measure squared distance from smooth version avg_edge_l = average_edge_length_by_vertex(predicted_mesh_verts, predicted_mesh_indices) avg_edge_l, _ = torch.broadcast_tensors(avg_edge_l.unsqueeze(1), lap) if self.edge_length_scale is not None: loss_per_vertex = (predicted_mesh_verts - lap) * self.edge_length_scale / avg_edge_l else: loss_per_vertex = predicted_mesh_verts - lap if self.scale_by_orig_smoothness: if self.orig is None: # Add a little bit so that nothing's / 0 self.orig = loss_per_vertex.clone().detach() + 0.01 # Scale by orig. If the orig loss for a given vertex was HIGH, # it's not so bad (divide) # if the orig loss for a given vertex was LOW, then it is bad! loss_per_vertex = loss_per_vertex - self.orig sqrt_error = torch.sqrt( torch.sum(loss_per_vertex * loss_per_vertex, dim=1)) mean_error = torch.mean(sqrt_error) return mean_error
def test(args, model, device, test_loader, epoch, writer): model.eval() test_loss = 0 correct = 0 isEval = False with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) data, _ = torch.broadcast_tensors( data, torch.zeros((steps, ) + data.shape)) data = data.permute(1, 2, 3, 4, 0) output = model(data) test_loss += F.cross_entropy( output, target, reduction='sum').item() # sum up batch loss pred = output.argmax( dim=1, keepdim=True) # get the index of the max log-probability correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) writer.add_scalar('Test Loss /epoch', test_loss, epoch) writer.add_scalar('Test Acc /epoch', 100. * correct / len(test_loader.dataset), epoch) for i, (name, param) in enumerate(model.named_parameters()): if '_s' in name: writer.add_histogram(name, param, epoch) print( '\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset)))
def broadcast_except(*tensors: torch.Tensor, dim=-1): shape = torch.broadcast_tensors(*[t.select(dim, 0) for t in tensors])[0].shape return [ t.expand(*shape[:t.ndim + dim + 1], t.shape[dim], *shape[t.ndim + dim + 1:]) for t in tensors ]
def broadcast_and_squeeze(*args): assert all([is_tensor(ar) for ar in args]), 'at least 1 object is not torch tensor' if all([np.prod(val.shape[2:]) == 1 for val in args]): args = [val.contiguous().view(size=val.shape[:2] + tuple([1, 1])) for val in args] uniformed_values = uniform_shapes(*args) broadcasted_values = torch.broadcast_tensors(*uniformed_values) return broadcasted_values
def _feature_dropout(self, p=0.5, training=True, inplace=False): """Randomly zeros out entire channels in the input tensor with probability :attr:`p`. (a channel is a nD feature map, e.g., the :math:`j`-th channel of the :math:`i`-th sample in the batched input is a nD tensor :math:`\text{input}[i, j]`).""" assert self.dim( ) >= 2, "feature dropout requires dimension to be at least 2" assert p >= 0.0 and p <= 1.0, "dropout probability has to be between 0 and 1" if training is False: if inplace: return self else: return self.clone() # take first 2 dimensions feature_dropout_size = self.size()[0:2] # create dropout tensor over the first two dimensions rand_tensor = crypten.mpc.rand(feature_dropout_size) feature_dropout_tensor = rand_tensor > p # Broadcast to remaining dimensions for i in range(2, self.dim()): feature_dropout_tensor = feature_dropout_tensor.unsqueeze(i) feature_dropout_tensor.share, self.share = torch.broadcast_tensors( feature_dropout_tensor.share, self.share) if inplace: result_tensor = self.mul_(feature_dropout_tensor).div_(1 - p) else: result_tensor = self.mul(feature_dropout_tensor).div_(1 - p) return result_tensor
def loss(self, item, pred): actual = item['pos'].unsqueeze(-2) mask = item['mask'].unsqueeze(-1) pred, actual = torch.broadcast_tensors(pred, actual) loss = F.mse_loss(pred, actual, reduction='none').sum(-1) loss *= mask return loss.mean()
def _grid_offsets(H, W, dtype, device): """ Returns the grid offsets HxWx2 within [-1,1] """ if (H, W) in VideoTools._offset_cache: return VideoTools._offset_cache[(H, W)] else: print("Create grid offsets for warping: W=%d, H=%d" % (W, H)) grid_offsetsH = torch.linspace(-1, +1, H, dtype=dtype, device=device) grid_offsetsW = torch.linspace(-1, +1, W, dtype=dtype, device=device) grid_offsetsH = torch.unsqueeze(grid_offsetsH, 1) grid_offsetsW = torch.unsqueeze(grid_offsetsW, 0) grid_offsets = torch.stack(torch.broadcast_tensors( grid_offsetsW, grid_offsetsH), dim=2) grid_offsets = torch.unsqueeze(grid_offsets, 0) # batch dimension grid_offsets = grid_offsets.detach() VideoTools._offset_cache[(H, W)] = grid_offsets return grid_offsets
def __init__(self, loc, precision_diag, belief, df, validate_args=True): precision_diag, belief, df = self.convert_float_params_to_tensor( loc, precision_diag, belief, df) if loc.dim() < 1 or precision_diag.dim() < 1 or df.dim( ) < 1 or belief.dim() < 1: raise ValueError( "loc, precision_diag, df, belief must be at least one-dimensional." ) if belief.size(-1) == 1 and precision_diag.size(-1) != 1: raise ValueError( "belief shouldn't end with dimensionality 1 if precision_diag doesn't" ) if df.size(-1) == 1 and precision_diag.size(-1) != 1: raise ValueError( "df shouldn't end with dimensionality 1 if precision_diag doesn't" ) df_, belief_ = df.unsqueeze(-1), belief.unsqueeze( -1) # add dim on right self.loc, self.precision_diag, df_, belief_ = torch.broadcast_tensors( loc, precision_diag, df_, belief_) self.df, self.belief = df_[..., 0], belief_[..., 0] # drop rightmost dim batch_shape, event_shape = self.loc.shape[:-1], self.loc.shape[-1:] self.dimensionality = event_shape.numel() if (self.df <= (self.dimensionality + 1)).any(): raise ValueError( "df must be greater than dimensionality + 1 to have expectation" ) super(NormalDiagonalWishart, self).__init__(batch_shape, event_shape, validate_args=validate_args)