Package: torchlie
torchlieA library for differentiable Lie groupsGetting StartedPrerequisitesWestronglyrecommend you install torchlie in a venv or conda environment with Python 3.8-3.10.torchlie requirestorchinstallation. To install for your particular CPU/CUDA configuration, follow the instructions in the PyTorchwebsite.InstallingpypipipinstalltorchlieFrom sourceThe simplest way to install torchlie from source is by running the followinggitclonehttps://github.com/facebookresearch/theseus.git&&cdtheseus/torchlie
pipinstall-e.If you are interested in contributing to torchlie, also installpipinstall-r../requirements/dev.txt
pre-commitinstalland follow the more detailed instructions inCONTRIBUTING.ExampleThe example below is also available as ascript.importtorchimporttorchlieaslieimporttorchlie.functionalaslieFbatch_size=5# ### Lie Tensor creation functionsg1=lie.SE3.rand(batch_size,requires_grad=True)print(f"Created SE3 tensor with shape{g1.shape}")g2=g1.clone()# Identity elementi1=lie.SO3.identity(2)i2=lie.SE3.identity(2)print("SO3 identity",i1,i1.shape)print("SE3 identity",i2,i2.shape)# Indexingg1_slice=g1[2:4]assertg1_slice.shape==(2,3,4)torch.testing.assert_close(g1_slice._t,g1._t[2:4])# type: ignoretry:bad=g1[3,2]exceptNotImplementedError:print("INDEXING ERROR: Can only slice the first dimension for now.")# ## Different constructorsg3_data=lieF.SO3.rand(5,requires_grad=True)# this is a regular tensor with SO3 data# Can create from a tensor as long as it's consistent with the desired ltypeg3=lie.from_tensor(g3_data,lie.SO3)# keeps grad historyassertg3.grad_fnisnotNonetry:x=lie.from_tensor(torch.zeros(1,3,3),lie.SO3)exceptValueErrorase:print(f"ERROR:{e}")defis_shared(t1,t2):# utility to check if memory is sharedreturnt1.storage().data_ptr()==t2.storage().data_ptr()# # Let's check different copy vs no-copy options# -- lie.SO3() lie.SE3()g3_leaf=lie.SO3(g3_data)# creates a leaf tensor and copies dataassertg3_leaf.grad_fnisNoneassertnotis_shared(g3_leaf,g3_data)# -- lie.LieTensor() constructor is equivalent to lie.SO3()g3_leaf_2=lie.LieTensor(g3_data,lie.SO3)assertg3_leaf_2.grad_fnisNoneassertnotis_shared(g3_leaf_2,g3_data)# -- as_lietensor()g4=lie.as_lietensor(g3_data,lie.SO3)assertis_shared(g3_data,g4)# shares storage if possibleassertg4.grad_fnisnotNone# result is not a leaf tensor# Calling with a LieTensor returns the same tensor...g5=lie.as_lietensor(g3,lie.SO3)assertg5isg3# ... unless dtype or device is differentg5_double=lie.as_lietensor(g3,lie.SO3,dtype=torch.double)assertg5_doubleisnotg3assertnotis_shared(g5_double,g3)# -- cast()g6=lie.cast(g3_data,lie.SO3)# alias for as_lietensor()assertis_shared(g3_data,g6)# -- LieTensor.new()g7=g3.new_lietensor(g3_data)assertnotis_shared(g3_data,g7)# doesn't share storageassertg7.grad_fnisNone# creates a leaf# ### Lie operationsv=torch.randn(batch_size,6)# Exponential and logarithmic mapout1=lie.SE3.exp(v)# also lie.exp(v, g1.ltype)print(f"Exp map returns a{type(out1)}.")out2=g1.log()# also lie.log(g1)print(f"Log map returns a{type(out2)}.")# Inverseout1=g1.inv()# also lie.inv(g1)# Compose# also lie.compose(g1, g2)out1=g1.compose(g2)# type: ignore# Differentiable jacobiansjacs,out=g1.jcompose(g2)# type: ignoreprint("Jacobians output is a 2-tuple.")print("    First element is a list of jacobians, one per group argument.")print(f"    For compose this means length{len(jacs)}.")print("    The second element of the tuple is the result of the operation itself.")print(f"    Which for compose is a{type(out).__name__}.")# Other options:#   * adj(), hat(), vee(), retract(), local(),#   * Jacobians: jlog(), jinv(), jexp()# ### Overriden operators# Composeout2=g1*g2torch.testing.assert_close(out1,out2,check_dtype=True)# Transfrom from (from local to world coordinate frame)p=torch.randn(batch_size,3)pt1=g1.transform_from(p)pt2=g1@ptorch.testing.assert_close(pt1,pt2)# For convenience, we provide a context to drop all ltype checks, and operate# on raw tensor data. However, keep in mind that this is prone to error.# Here is one example of how this works.withlie.as_euclidean():gg1=torch.sin(g1)# The above is the same as this next call, but the context might be more convenient# if one is doing similar hacky stuff on several group objects.gg2=torch.sin(g1._t)torch.testing.assert_close(gg1,gg2)print("Success: We just did some ops that make no sense for SE3 tensors.")# ### Lie tensors can also be used as leaf tensors for torch optimizersg1=lie.SE3.rand(1,requires_grad=True)g2=lie.SE3.rand(1)opt=torch.optim.Adam([g1],lr=0.1)foriinrange(10):opt.zero_grad()d=g1.local(g2)loss=torch.sum(d**2)loss.backward()opt.step()print(f"Iter{i}. Loss:{loss.item():.3f}")Citing torchlieIf you use torchlie in your work, please cite thepaperwith the BibTeX below.@article{pineda2022theseus,title={{Theseus: A Library for Differentiable Nonlinear Optimization}},author={Luis Pineda and Taosha Fan and Maurizio Monge and Shobha Venkataraman and Paloma Sodhi and Ricky TQ Chen and Joseph Ortiz and Daniel DeTone and Austin Wang and Stuart Anderson and Jing Dong and Brandon Amos and Mustafa Mukadam},journal={Advances in Neural Information Processing Systems},year={2022}}Licensetorchlie is MIT licensed. See theLICENSEfor details.Additional InformationJoin the community onGithub Discussionsfor questions and sugesstions.UseGithub Issuesfor bugs and features.SeeCONTRIBUTINGif interested in helping out.