[Feat] Improve device switching for transforms

🚀 Feature

Improve device switching for transforms

Motivation

Currently tensors are stored as is within transforms (I’m using kornia.Rotate as example here, but this same is valid for many other transforms). This, makes the switching between devices inconvenient, since torch.nn.Module.to() does not affect them:

import torch
import kornia

angle = torch.tensor(30.0).view(1, -1)

transform = kornia.Rotate(angle)
print(transform.angle.device)

transform = transform.to("cuda")
print(transform.angle.device)

cpu
cpu

Pitch

Register all tensors as buffers:

from torch import nn


class RotateMock(nn.Module):
    def __init__(self, angle: torch.Tensor) -> None:
        super().__init__()
        self.register_buffer("angle", angle)

transform = RotateMock(angle)
print(transform.angle.device)

transform = transform.to("cuda")
print(transform.angle.device)

cpu
cuda:0

[pmeier]

Nope. I’ll see if I have time this week to send a PR. Otherwise I’ll get back to you.

[shijianjian]

It shall be enough.

import torch
from torch import nn


class RotateMock(nn.Module):
    def __init__(self, angle: torch.Tensor) -> None:
        super().__init__()
        self.angle = nn.Parameter(angle)
    def forward(self, input):
        return input * self.angle

r = RotateMock(torch.tensor(0.5, requires_grad=True))
res = r(torch.tensor(2.))

criterion = torch.nn.L1Loss() 
optimizer = torch.optim.SGD(r.parameters(), lr=10)

loss = criterion(res, torch.tensor(5.))

loss.backward()
optimizer.step()

print(r.angle, r.angle.grad)

Parameter containing:
tensor(20.5000, requires_grad=True) tensor(-2.)