What's Changed
* dev(hansbug): try adapt torch2 by HansBug in https://github.com/opendilab/treevalue/pull/85
* dev(narugo): torch compile's test by HansBug in https://github.com/opendilab/treevalue/pull/87
In the new version (v1.4.12), support for torch >= 2 versions has been added, including support for `torch.compile` for faster inference and backpropagation. Here's an example:
python
from typing import Tuple, Mapping
import torch
from torch import nn
from treevalue import FastTreeValue
A simple MLP
class MLP(nn.Module):
def __init__(self, in_features: int, out_features: int, layers: Tuple[int, ...] = (1024,)):
nn.Module.__init__(self)
self.in_features = in_features
self.out_features = out_features
self.layers = layers
ios = [self.in_features, *self.layers, self.out_features]
self.mlp = nn.Sequential(
*(
nn.Linear(in_, out_, bias=True)
for in_, out_ in zip(ios[:-1], ios[1:])
)
)
def forward(self, x):
return self.mlp(x)
Multiple headed MLP
class MultiHeadMLP(nn.Module):
def __init__(self, in_features: int, out_features: Mapping[str, int], layers: Tuple[int, ...] = (1024,)):
nn.Module.__init__(self)
self.in_features = in_features
self.out_features = out_features
self.layers = layers
_networks = {
o_name: MLP(in_features, o_feat, layers)
for o_name, o_feat in self.out_features.items()
}
self.mlps = nn.ModuleDict(_networks) use nn.ModuleDict to register child MLPs
self._t_mlps = FastTreeValue(_networks) use TreeValue for batch inferring
def forward(self, x):
return self._t_mlps(x)
if __name__ == '__main__':
net = MultiHeadMLP(
20,
{'a': 10, 'b': 20, 'c': 14, 'd': 3},
)
net = torch.compile(net)
print(net)
input_ = torch.randn(1, 10, 20)
output = net(input_)
print(output.shape)
The compiled version of the MultiHeadMLP above will have the following network structure:
text
OptimizedModule(
(_orig_mod): MultiHeadMLP(
(mlps): ModuleDict(
(a): MLP(
(mlp): Sequential(
(0): Linear(in_features=20, out_features=1024, bias=True)
(1): Linear(in_features=1024, out_features=10, bias=True)
)
)
(b): MLP(
(mlp): Sequential(
(0): Linear(in_features=20, out_features=1024, bias=True)
(1): Linear(in_features=1024, out_features=20, bias=True)
)
)
(c): MLP(
(mlp): Sequential(
(0): Linear(in_features=20, out_features=1024, bias=True)
(1): Linear(in_features=1024, out_features=14, bias=True)
)
)
(d): MLP(
(mlp): Sequential(
(0): Linear(in_features=20, out_features=1024, bias=True)
(1): Linear(in_features=1024, out_features=3, bias=True)
)
)
)
)
)
And the inference output after passing `float32[1, 10, 20]` as input will have the following dimensions:
text
<FastTreeValue 0x7fe9b197e6a0>
├── 'a' --> torch.Size([1, 10, 10])
├── 'b' --> torch.Size([1, 10, 20])
├── 'c' --> torch.Size([1, 10, 14])
└── 'd' --> torch.Size([1, 10, 3])
**Full Changelog**: https://github.com/opendilab/treevalue/compare/v1.4.11...v1.4.12