64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def unsqueeze(input, dims):
|
|
""" Implement multi-dimension unsqueeze function. """
|
|
if isinstance(dims, (list, tuple)):
|
|
dims = [
|
|
dim if dim >= 0 else dim + len(input.shape) + 1 for dim in dims
|
|
]
|
|
dims = sorted(dims, reverse=True)
|
|
shape = list(input.shape)
|
|
for dim in dims:
|
|
shape.insert(dim, 1)
|
|
return torch.reshape(input, shape)
|
|
elif isinstance(dims, int):
|
|
return input.unsqueeze(dims)
|
|
else:
|
|
raise ValueError('Warning: type(dims) must in (list, tuple, int)!')
|
|
|
|
|
|
def gumbel_softmax(input, tau=1, eps=1e-10):
|
|
""" Basic implement of gumbel_softmax. """
|
|
U = torch.tensor(np.random.rand(*input.shape))
|
|
gumbel = 0.0 - torch.log(eps - torch.log(U + eps))
|
|
y = input + gumbel
|
|
return F.softmax(y / tau)
|
|
|
|
|
|
def equal(x, y, dtype=None):
|
|
""" Implement equal in dygraph mode. (paddle) """
|
|
if dtype is None:
|
|
dtype = 'float32'
|
|
if isinstance(x, torch.Tensor):
|
|
x = x.numpy()
|
|
if isinstance(y, torch.Tensor):
|
|
y = y.numpy()
|
|
out = np.equal(x, y).astype(dtype)
|
|
return torch.tensor(out)
|
|
|
|
|
|
def not_equal(x, y, dtype=None):
|
|
""" Implement not_equal in dygraph mode. (paddle) """
|
|
return 1 - equal(x, y, dtype)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
a = torch.tensor([[1, 1], [3, 4]])
|
|
b = torch.tensor([[1, 1], [3, 4]])
|
|
c = torch.equal(a, a)
|
|
c1 = equal(a, 3)
|
|
d = 1 - torch.not_equal(a, 3).float()
|
|
print(c)
|
|
print(c1)
|
|
print(d)
|
|
e = F.gumbel_softmax(a)
|
|
f = a.unsqueeze(a)
|
|
g = unsqueeze(a, dims=[0, 0, 1])
|
|
print(g, g.shape)
|