pytorch-extras
pip install pytorch-extras
Usage
expand_along
expand_along(var, mask) - Useful for selecting a dynamic amount of items from different indexes using a byte mask. This is a bit like numpy.repeat.
import torch
import torch_extras
setattr(torch, 'expand_along', torch_extras.expand_along)
var = torch.Tensor([1, 0, 2])
mask = torch.ByteTensor([[True, True], [False, True], [False, False]])
torch.expand_along(var, mask)
# (1, 1, 0)
expand_dims
expand_dims(var, dim) - Is similar to numpy.expand_dims.
import torch
import torch_extras
setattr(torch, 'expand_dims', torch_extras.expand_dims)
var = torch.range(0, 9).view(-1, 2)
torch.expand_dims(var, 0).size()
# (1, 5, 2)
Note: Have recently found out about torch.unsqeeze, which has the same API and is probably a more effective method for expanding dimensions.
select_item
select_item(var, index) - Is similar to [var[row,col] for row, col in enumerate(index)].
import torch
import torch_extras
setattr(torch, 'select_item', torch_extras.select_item)
var = torch.range(0, 9).view(-1, 2)
index = torch.LongTensor([0, 0, 0, 1, 1])
torch.select_item(var, index)
# [0, 2, 4, 7, 9]
cast
cast(var, type) - Cast a Tensor to the given type.
import torch
import torch_extras
setattr(torch, 'cast', torch_extras.cast)
input = torch.FloatTensor(1)
target_type = type(torch.LongTensor(1))
type(torch.cast(input, target_type))
# <class 'torch.LongTensor'>
one_hot
one_hot(size, index) - Creates a matrix of one hot vectors.
import torch
import torch_extras
setattr(torch, 'one_hot', torch_extras.one_hot)
size = (3, 3)
index = torch.LongTensor([2, 0, 1]).view(-1, 1)
torch.one_hot(size, index)
# [[0, 0, 1], [1, 0, 0], [0, 1, 0]]
nll
nll(log_prob, label) - Is similar to nll_loss except does not return an aggregate.
import torch
from torch.autograd import Variable
import torch.nn.functional as F
import torch_extras
setattr(torch, 'nll', torch_extras.nll)
input = Variable(torch.FloatTensor([[0.5, 0.2, 0.3], [0.1, 0.8, 0.1]]))
target = Variable(torch.LongTensor([1, 2]).view(-1, 1))
output = torch.nll(torch.log(input), target)
output.size()
# (2, 1)