functorch
functorch -> torch.func (as of PyTorch 2.0)
Por que functorch? | Guia de instalação | Transformações | Documentação | Planos futuros
Atualmente, esta biblioteca está em desenvolvimento pesado - se você tiver sugestões na API ou em casos de uso que gostaria de ser coberto, abra um problema do GitHub ou alcance. Gostaríamos muito de saber como você está usando a biblioteca.
functorch são transformadas de função composta do tipo JAX para Pytorch.
O objetivo é fornecer transformados composíveis de vmap e grad que funcionam com módulos Pytorch e Pytorch AutoGRAD com bom desempenho no modo ansioso.
Além disso, existe uma funcionalidade experimental para rastrear essas transformações usando FX para capturar os resultados dessas transformações antes do tempo. Isso nos permitiria compilar os resultados do VMAP ou GRAD para melhorar o desempenho.
Hoje existem vários casos de uso que são complicados em Pytorch:
A composição de transformações vmap , grad , vjp e jvp nos permite expressar o acima, sem projetar um subsistema separado para cada um. Essa idéia de transformadas de função composta vem da estrutura JAX.
Existem duas maneiras de instalar o Functorch:
Recomendamos experimentar o Functorch Beta primeiro.
Siga as instruções neste caderno Colab
Em 21/09/2022, functorch vem instalado ao lado de um binário noturno de Pytorch. Por favor, instale uma visualização (noturna) pytorch binária; Consulte https://pytorch.org/ para obter instruções.
Depois de fazer isso, execute uma rápida verificação de sanidade no Python:
import torch
from functorch import vmap
x = torch . randn ( 3 )
y = vmap ( torch . sin )( x )
assert torch . allclose ( y , x . sin ()) A partir de 21/09/2022, functorch vem instalado ao lado de Pytorch e está na árvore de origem Pytorch. Por favor, instale o Pytorch da fonte, então você poderá import functorch .
Tente executar alguns testes para garantir que tudo esteja bem:
pytest test/test_vmap.py -v
pytest test/test_eager_transforms.py -vAoTaUTograd possui alguns requisitos opcionais adicionais. Você pode instalá -los via:
pip install networkx Para executar testes Functorch, instale nossas dependências de teste ( expecttest , pyyaml ).
Siga as instruções aqui
Pré -requisito: Instale Pytorch
pip install functorchFinalmente, execute uma rápida verificação de sanidade em Python:
import torch
from functorch import vmap
x = torch . randn ( 3 )
y = vmap ( torch . sin )( x )
assert torch . allclose ( y , x . sin ())No momento, apoiamos as seguintes transformações:
grad , vjp , jvp ,jacrev , jacfwd , hessianvmapAlém disso, temos alguns utilitários para trabalhar com módulos Pytorch.
make_functional(model)make_functional_with_buffers(model) Nota: vmap impõe restrições ao código em que pode ser usado. Para mais detalhes, leia o seu Docstring.
vmap(func)(*inputs) é uma transformação que adiciona uma dimensão a todas as operações tensoras no func . vmap(func) retorna uma nova função que mapeia func em alguma dimensão (padrão: 0) de cada tensor nas inputs .
vmap é útil para esconder as dimensões do lote: pode -se escrever uma função func que é executada em exemplos e, em seguida, eleva -a para uma função que pode levar lotes de exemplos com vmap(func) , levando a uma experiência de modelagem mais simples:
from functorch import vmap
batch_size , feature_size = 3 , 5
weights = torch . randn ( feature_size , requires_grad = True )
def model ( feature_vec ):
# Very simple linear model with activation
assert feature_vec . dim () == 1
return feature_vec . dot ( weights ). relu ()
examples = torch . randn ( batch_size , feature_size )
result = vmap ( model )( examples ) grad(func)(*inputs) assume que func retorna um tensor de elemento único. Ele calcula os gradientes da saída do func WRT para inputs[0] .
from functorch import grad
x = torch . randn ([])
cos_x = grad ( lambda x : torch . sin ( x ))( x )
assert torch . allclose ( cos_x , x . cos ())
# Second-order gradients
neg_sin_x = grad ( grad ( lambda x : torch . sin ( x )))( x )
assert torch . allclose ( neg_sin_x , - x . sin ()) Quando composto com vmap , grad pode ser usado para calcular gradientes por amostra:
from functorch import vmap
batch_size , feature_size = 3 , 5
def model ( weights , feature_vec ):
# Very simple linear model with activation
assert feature_vec . dim () == 1
return feature_vec . dot ( weights ). relu ()
def compute_loss ( weights , example , target ):
y = model ( weights , example )
return (( y - target ) ** 2 ). mean () # MSELoss
weights = torch . randn ( feature_size , requires_grad = True )
examples = torch . randn ( batch_size , feature_size )
targets = torch . randn ( batch_size )
inputs = ( weights , examples , targets )
grad_weight_per_example = vmap ( grad ( compute_loss ), in_dims = ( None , 0 , 0 ))( * inputs ) A transformação vjp aplica func às inputs e retorna uma nova função que calcula os VJPs, com alguns tensores cotangents .
from functorch import vjp
outputs , vjp_fn = vjp ( func , inputs ); vjps = vjp_fn ( * cotangents ) A jvp transforma calcula os produtos de vetor jacobiano e também é conhecido como "AD para o modo avançado". Não é uma função de ordem superior, diferentemente da maioria das outras transformações, mas retorna as saídas do func(inputs) , bem como as jvp .
from functorch import jvp
x = torch . randn ( 5 )
y = torch . randn ( 5 )
f = lambda x , y : ( x * y )
_ , output = jvp ( f , ( x , y ), ( torch . ones ( 5 ), torch . ones ( 5 )))
assert torch . allclose ( output , x + y ) A transformação jacrev retorna uma nova função que recebe x e retorna o jacobiano de torch.sin em relação a x usando anúncio no modo reverso.
from functorch import jacrev
x = torch . randn ( 5 )
jacobian = jacrev ( torch . sin )( x )
expected = torch . diag ( torch . cos ( x ))
assert torch . allclose ( jacobian , expected ) Use jacrev para calcular o jacobiano. Isso pode ser composto com o VMAP para produzir jacobianos em lotes:
x = torch . randn ( 64 , 5 )
jacobian = vmap ( jacrev ( torch . sin ))( x )
assert jacobian . shape == ( 64 , 5 , 5 ) jacfwd é um substituto para jacrev que calcula os jacobianos usando o anúncio de modo avançado:
from functorch import jacfwd
x = torch . randn ( 5 )
jacobian = jacfwd ( torch . sin )( x )
expected = torch . diag ( torch . cos ( x ))
assert torch . allclose ( jacobian , expected ) Compondo jacrev consigo mesmo ou jacfwd pode produzir hessianos:
def f ( x ):
return x . sin (). sum ()
x = torch . randn ( 5 )
hessian0 = jacrev ( jacrev ( f ))( x )
hessian1 = jacfwd ( jacrev ( f ))( x ) O hessian é uma função de conveniência que combina jacfwd e jacrev :
from functorch import hessian
def f ( x ):
return x . sin (). sum ()
x = torch . randn ( 5 )
hess = hessian ( f )( x ) Também podemos rastrear essas transformações para capturar os resultados como novo código usando make_fx . Há também integração experimental com o compilador NNC (funciona apenas na CPU por enquanto!).
from functorch import make_fx , grad
def f ( x ):
return torch . sin ( x ). sum ()
x = torch . randn ( 100 )
grad_f = make_fx ( grad ( f ))( x )
print ( grad_f . code )
def forward ( self , x_1 ):
sin = torch . ops . aten . sin ( x_1 )
sum_1 = torch . ops . aten . sum ( sin , None ); sin = None
cos = torch . ops . aten . cos ( x_1 ); x_1 = None
_tensor_constant0 = self . _tensor_constant0
mul = torch . ops . aten . mul ( _tensor_constant0 , cos ); _tensor_constant0 = cos = None
return mulÀs vezes, você pode realizar uma transformação em relação aos parâmetros e/ou buffers de um nn.module. Isso pode acontecer, por exemplo, em:
Nossa solução para isso agora é uma API que, dada a Nn.Module, cria uma versão sem estado que pode ser chamada de função.
make_functional(model) retorna uma versão funcional do model e o model.parameters()make_functional_with_buffers(model) retorna uma versão funcional do model e o model.parameters() e model.buffers() .Aqui está um exemplo em que calculamos os gradientes por amostra usando uma camada nn.Linear:
import torch
from functorch import make_functional , vmap , grad
model = torch . nn . Linear ( 3 , 3 )
data = torch . randn ( 64 , 3 )
targets = torch . randn ( 64 , 3 )
func_model , params = make_functional ( model )
def compute_loss ( params , data , targets ):
preds = func_model ( params , data )
return torch . mean (( preds - targets ) ** 2 )
per_sample_grads = vmap ( grad ( compute_loss ), ( None , 0 , 0 ))( params , data , targets ) Se você estiver fazendo um conjunto de modelos, poderá encontrar combine_state_for_ensemble útil.
Para mais documentação, consulte o site do nosso Docs.
torch._C._functorch.dump_tensor : DUMPS DISPAÇÃO TECHAS NA STACK torch._C._functorch._set_vmap_fallback_warning_enabled(False) Se o spam de aviso do VMAP o incomodar.
No estado final, gostaríamos de usar isso em Pytorch depois que resolvemos os detalhes do projeto. Para descobrir os detalhes, precisamos da sua ajuda - envie -nos seus casos de uso iniciando uma conversa no rastreador de questões ou tentando nosso projeto.
O Functorch possui uma licença de estilo BSD, conforme encontrado no arquivo de licença.
Se você usar o Functorch em sua publicação, cite -o usando a seguinte entrada Bibtex.
@Misc { functorch2021 ,
author = { Horace He, Richard Zou } ,
title = { functorch: JAX-like composable function transforms for PyTorch } ,
howpublished = { url{https://github.com/pytorch/functorch} } ,
year = { 2021 }
}
