| 7 | |
| 8 | |
| 9 | class CM(autograd.Function): |
| 10 | |
| 11 | @staticmethod |
| 12 | def forward(ctx, inputs, targets, features, momentum): |
| 13 | ctx.features = features |
| 14 | ctx.momentum = momentum |
| 15 | ctx.save_for_backward(inputs, targets) |
| 16 | outputs = inputs.mm(ctx.features.t()) |
| 17 | |
| 18 | return outputs |
| 19 | |
| 20 | @staticmethod |
| 21 | def backward(ctx, grad_outputs): |
| 22 | inputs, targets = ctx.saved_tensors |
| 23 | grad_inputs = None |
| 24 | if ctx.needs_input_grad[0]: |
| 25 | grad_inputs = grad_outputs.mm(ctx.features) |
| 26 | |
| 27 | # momentum update |
| 28 | for x, y in zip(inputs, targets): |
| 29 | ctx.features[y] = ctx.momentum * ctx.features[y] + (1. - ctx.momentum) * x |
| 30 | ctx.features[y] /= ctx.features[y].norm() |
| 31 | |
| 32 | return grad_inputs, None, None, None |
| 33 | |
| 34 | |
| 35 | def cm(inputs, indexes, features, momentum=0.5, device=None): |
nothing calls this directly
no outgoing calls
no test coverage detected