【Pytorch】便利なやつら

はじめに

Pytorchで知らない関数があったので備忘録として残します。

拡散モデルでよくつかうやつら

torch.cumprod

$y=x_1\times x_2\times\cdots x_N$

>> import torch
>> a = torch.tensor([1,2,3])
>> torch.cumprod(a,dim=0)
tensor([1,2,6])

self.register_buffer

パラメータ$a$は最適化されるが、パラメータ$b$は固定される。

import torch

class MyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.a = nn.Parameter(torch.randn(10))
        self.register_buffer('b', torch.randn(10))

    def forward(self, x):
        return self.a*x + b

torch.clamp

>> import torch
>> a = torch.randn(4)
>> a
tensor([-1.7120,  0.1734, -0.0478, -0.0922])
>> torch.clamp(a, min=-0.5, max=0.5)
tensor([-0.5000,  0.1734, -0.0478, -0.0922])

torch.rand_like

>> import torch
>> noise = torch.randn_like(x)

torch.from_numpy

>> import torch
>> import numpy as np
>> a = np.array([1, 2, 3])
>> t = torch.from_numpy(a).float()
tensor([1.,2.,3.])

ニューラルネットワークのパラメータの更新

def soft_update_from_to(source, target, tau):
    for target_param, param in zip(target.parameters(), source.parameters()):
        target_param.data.copy_(
            target_param.data * (1.0 - tau) + param.data * tau
        )