8000 [primTorch] Implement NLL loss reference by rdspring1 · Pull Request #81128 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

[primTorch] Implement NLL loss reference #81128

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 29 commits into from
Closed

[primTorch] Implement NLL loss reference #81128

Changes from 1 commit
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
2256d36
Initial nll_loss implementation
rdspring1 Jun 29, 2022
b1393e5
fixup
rdspring1 Jun 29, 2022
f72db25
Disable validate_view_consistency check
rdspring1 Jun 29, 2022
055e0e2
Merge 1d and 2d nll_loss functions
rdspring1 Jun 29, 2022
96cc303
Add target class check - disabled because of FakeTensor
rdspring1 Jun 29, 2022
370bc60
refactor helper function
rdspring1 Jul 8, 2022
612ce91
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Sep 25, 2022
e7a3ae4
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Sep 27, 2022
44702b8
Address comments - rnd 1
rdspring1 Sep 27, 2022
c71d746
fixup
rdspring1 Sep 27, 2022
e0554f2
Refactor class weight selection
rdspring1 Sep 28, 2022
6aa6b62
Add comments
rdspring1 Sep 28, 2022
dde53e3
Replace 4-D case for image inputs with general 3-D case
rdspring1 Sep 28, 2022
4df9971
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Sep 28, 2022
39883b6
add comments
rdspring1 Sep 28, 2022
1a635cd
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Sep 28, 2022
590866b
8000 Add class check
rdspring1 Sep 28, 2022
1b88f57
Add FakeTensor Issue
rdspring1 Sep 28, 2022
c59279e
add zero-dim check
rdspring1 Sep 28, 2022
e6d01e4
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Sep 30, 2022
f2c9c3f
Update comments
rdspring1 Sep 30, 2022
10b85ff
fixup
rdspring1 Sep 30, 2022
96a6142
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Oct 3, 2022
6cbdf01
lint
rdspring1 Oct 3, 2022
746a60e
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Oct 11, 2022
e1eb641
Merge branch 'master' of github.com:rdspring1/pytorch into ref_nll_loss
rdspring1 Oct 16, 2022
ef5719e
PR comments
rdspring1 Oct 16, 2022
76bfc80
update test args
rdspring1 Oct 16, 2022
3cd82ab
add type promotion wrapper
rdspring1 Oct 17, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Merge 1d and 2d nll_loss functions
  • Loading branch information
@rdspring1
rdspring1 committed Jul 8, 2022
commit 055e0e2094131d8c5d7474ece49c5c27ab44260c
90 changes: 22 additions & 68 deletions torch/_refs/nn/functional/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -366,22 +366,18 @@ def _get_string_reduction_arg(
return ret


def _nll_loss_2d(
def _nll_loss_nd(
input: TensorLikeType,
target: TensorLikeType,
weight: Optional[TensorLikeType],
reduction: str,
ignore_index: int,
) -> TensorLikeType:
if input.ndim != 4:
msg = "Expected 4 dimensions for input but got {}."
if input 8000 .ndim == 3 or input.ndim > 4:
msg = "Expected input dimension to be either [1, 2, 4] but recieved {}."
raise ValueError(msg.format(input.ndim))

if target.ndim != 3:
msg = "Expected 3 dimensions for input but got {}."
raise ValueError(msg.format(input.ndim))

if input.shape[0] != target.shape[0]:
if input.ndim != 1 and input.shape[0] != target.shape[0]:
msg = "Expected input batch size ({}) to match target batch size ({})."
raise ValueError(msg.format(input.shape[0], target.shape[0]))

Expand All @@ -403,62 +399,22 @@ def _nll_loss_2d(
weight[flat_target],
)

batch_size = input.shape[0]
height = input.shape[2]
width = input.shape[3]
extent = height * width
numel = batch_size * extent

bdx = torch.arange(numel) // extent
hdx = (torch.arange(numel) - (bdx * extent)) // width
wdx = torch.arange(numel) % width

loss = -input[bdx, flat_target, hdx, wdx] * current_weight
loss = torch.reshape(loss, target.shape)

if reduction == "none":
return loss
elif reduction == "sum":
return torch.sum(loss)
else:
return torch.sum(loss) / torch.sum(current_weight)


def _nll_loss_1d(
input: TensorLikeType,
target: TensorLikeType,
weight: Optional[TensorLikeType],
reduction: str,
ignore_index: int,
) -> TensorLikeType:
if input.ndim < 1 or input.ndim > 2:
8000 msg = "Expected 1 or 2 dimensions for input but got {}."
raise ValueError(msg.format(input.ndim))

if input.ndim != 1 and input.shape[0] != target.shape[0]:
msg = "Expected input batch size ({}) to match target batch size ({})."
raise ValueError(msg.format(input.shape[0], target.shape[0]))

_check_reduction_value(reduction)

ignore_classes_mask = torch.eq(target, ignore_index)
ignore_class_weight = torch.scalar_tensor(0, dtype=input.dtype, device=input.device)
if weight is None:
current_weight = torch.where(
ignore_classes_mask,
ignore_class_weight,
torch.scalar_tensor(1, dtype=input.dtype, device=input.device),
)
else:
current_weight = torch.where(
ignore_classes_mask, ignore_class_weight.expand_as(target), weight[target]
)

if input.ndim == 1:
loss = -input[target] * current_weight
else:
elif input.ndim == 2:
batch_size = input.shape[0]
loss = -input[torch.arange(batch_size), target] * current_weight
else:
batch_size = input.shape[0]
height = input.shape[2]
width = input.shape[3]
extent = height * width
numel = batch_size * extent
bdx = torch.arange(numel) // extent
hdx = (torch.arange(numel) - (bdx * extent)) // width
wdx = torch.arange(numel) % width
loss = -input[bdx, flat_target, hdx, wdx] * current_weight
loss = torch.reshape(loss, target.shape)

if reduction == "none":
return loss
Expand All @@ -478,15 +434,11 @@ def nll_loss(
reduction: str = "mean",
) -> TensorLikeType:
if size_average is not None or reduce is not None:
# TODO raise exception instead of converting value
# TODO: raise exception instead of converting value
# msg = "size_average and reduce args are deprecated, please use reduction argument."
reduction = _get_string_reduction_arg(size_average, reduce)

if input.ndim == 1 or input.ndim == 2:
return _nll_loss_1d(input, target, weight, reduction, ignore_index)
elif input.ndim == 4:
return _nll_loss_2d(input, target, weight, reduction, ignore_index)
else:
if input.ndim == 3 or input.ndim > 4:
# input ndim is == 3 or > 4
batch_size = input.shape[0]
num_classes = input.shape[1]
Expand All @@ -508,10 +460,12 @@ def nll_loss(
target = torch.reshape(target, [batch_size, 0, 0])

if reduction == "none":
return _nll_loss_2d(input, target, weight, reduction, ignore_index)
return _nll_loss_nd(input, target, weight, reduction, ignore_index)
else:
result = _nll_loss_2d(input, target, weight, reduction, ignore_index)
result = _nll_loss_nd(input, target, weight, reduction, ignore_index)
return torch.reshape(result, out_size)
else:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a comment describing what this else branch is for.

From a code organization and readability standpoing these branches seem a little odd. Maybe we can explain them better?

In particular -- can input be zero or one dimension? If so, how do we interpret that? The documentation for suggests that input should have at least two dimensions. And why are inputs with three or four dimensions special?

Finally, prefer putting shorter branches which short-circuit first. That typically lets code have fewer indentation levels:

# shortcircuits if foo because...
if foo:
  return x

# implicit else branch doesn't have to be indented
...

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I refactored _nll_loss_nd to handle 1-3 dimensions. If there are more than 3 dimensions, the k-dimension is flattened to create a 3D tensor. The Aten implementation used a 4D case for image inputs.

    # The _nll_loss_nd helper function handles the most common cases.
    # ndim == 1 (Single Example)
    #   => Batch Size: 1, Input: (C), Target: ()
    # ndim == 2 (k = 1)
    #   => Batch Size: N, Input: (N, C), Target: (N)
    # ndim == 3 (k > 1)
    #   => Batch Size: N, Input: (N, C, K), Target: (N, K)
    # ndim > 3
    #   => reshape the input and target to the 3-D case

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is the 4D case interesting to model here?

return _nll_loss_nd(input, target, weight, reduction, ignore_index)


# tanhshrink does not use _make_elementwise_unary_reference because it does not support out
Expand Down
0