PyTorch: merge layer with a constant input
Closed this issue · 15 comments
Prerequisites
Please make sure to check off these prerequisites before submitting a bug report.
- Test that the bug appears on the current version of the master branch. Make sure to include the commit hash of the commit you checked out.
- Check that the issue hasn't already been reported, by checking the currently open issues.
- If there are steps to reproduce the problem, make sure to write them down below.
- If relevant, please include the hls4ml project files, which were created directly before and/or after the bug.
Quick summary
Merge layers that have a constant input do not work with PyTorch.
Details
Steps to Reproduce
- Clone the hls4ml repository
- Checkout the master branch, with commit hash: [afed23b]
- Run this test code.
from pathlib import Path
import numpy as np
import os
import shutil
import torch
import torch.nn as nn
from torchinfo import summary
from hls4ml.converters import convert_from_pytorch_model
from hls4ml.utils.config import config_from_pytorch_model
test_root_path = Path(__file__).parent
if __name__ == "__main__":
class test(nn.Module):
def __init__(self):
super().__init__()
self.downsample = nn.AvgPool1d(kernel_size=1, stride=2)
def forward(self, x):
d = self.downsample(x)
p = torch.mul(d,4.3)
return torch.cat((d, p), dim=-1)
n_in = 2
size_in = 8
n_batch = 2
model = test()
io_type='io_stream'
backend='Vitis'
output_dir = str(test_root_path / f'hls4mlprj_mul_{backend}_{io_type}')
if os.path.exists(output_dir):
print("delete project dir")
shutil.rmtree(output_dir)
model.eval()
summary(model, input_size=(n_batch, n_in, size_in))
X_input = np.random.rand(n_batch, n_in, size_in)
with torch.no_grad():
pytorch_prediction = model(torch.Tensor(X_input)).detach().numpy()
# X_input is channels last
X_input_hls = np.ascontiguousarray(X_input.transpose(0, 2, 1))
# write tb data
ipf = "./tb_input_features.dat"
if os.path.isfile(ipf):
os.remove(ipf)
np.savetxt(ipf, X_input_hls.flatten(), newline=" ")
opf = "./tb_output_predictions.dat"
if os.path.isfile(opf):
os.remove(opf)
with open(opf, "ab") as f:
for p in pytorch_prediction:
np.savetxt(f, p.flatten(), newline=" ")
config = config_from_pytorch_model(model,
(None, n_in, size_in),
backend=backend,
default_precision='ap_fixed<16,6>',
channels_last_conversion='internal',
transpose_outputs=False)
config['Model']['Strategy'] = 'Resource'
print(config)
print(output_dir)
hls_model = convert_from_pytorch_model(
model,
output_dir=output_dir,
input_data_tb=ipf,
output_data_tb=opf,
backend=backend,
hls_config=config,
io_type=io_type,
part='xcvu9p-flga2104-2-e'
)
hls_model.compile()
print("pytorch_prediction")
print(pytorch_prediction)
print("pytorch_prediction.shape: ", end=" ")
print(pytorch_prediction.shape)
# reshape hls prediction to channels last, then transpose, then reshape
# to match .view
hls_prediction = hls_model.predict(X_input_hls)
#hls_prediction = np.transpose(
# np.reshape(hls_prediction,
# (n_batch, int(size_in/2)+size_in, n_out)),
# (0,2,1)
#)
print("hls_prediction")
print(hls_prediction)
print("hls_prediction.shape: ", end=" ")
print(hls_prediction.shape)
rtol = 1.0e-2
atol = 1.0e-2
assert len(pytorch_prediction) == len(hls_prediction), "length mismatch"
assert pytorch_prediction.shape == hls_prediction.shape, "shape mismatch"
for p, h in zip(pytorch_prediction, hls_prediction):
np.testing.assert_allclose(p,
h,
rtol=rtol, atol=atol)
# synthesize
hls_model.build(csim=True, synth=True, cosim=True, validation=True)
Expected behavior
Successful synthesis.
Actual behavior
==========================================================================================
Layer (type:depth-idx) Output Shape Param #
==========================================================================================
test [2, 2, 8] 12
├─AvgPool1d: 1-1 [2, 2, 4] --
==========================================================================================
Total params: 12
Trainable params: 12
Non-trainable params: 0
Total mult-adds (M): 0
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.00
Params size (MB): 0.00
Estimated Total Size (MB): 0.00
==========================================================================================
{'Model': {'Precision': 'ap_fixed<16,6>', 'ReuseFactor': 1, 'ChannelsLastConversion': 'internal', 'TransposeOutputs': False, 'Strategy': 'Resource'}, 'PytorchModel': test(
(conv1): Conv1d(2, 2, kernel_size=(3,), stride=(1,), padding=(1,), bias=False)
(downsample): AvgPool1d(kernel_size=(1,), stride=(2,), padding=(0,))
), 'InputShape': (None, 2, 8)}
/home/hls4ml-user/work/ewstapp_research/isolate/NETWORK/hls4mlprj_mul_Vitis_io_stream
Interpreting Model ...
Topology:
Layer name: downsample, layer type: AveragePooling1D, input shape: [[None, 2, 8]]
Layer name: mul, layer type: Merge, input shape: [[None, 2, 4]]
Layer name: cat, layer type: Concatenate, input shape: [[None, 2, 4], [None, 2, 4]]
Creating HLS model
WARNING: Changing pipeline style to "dataflow".
Traceback (most recent call last):
File "/home/hls4ml-user/work/ewstapp_research/isolate/NETWORK/test_mul.py", line 81, in <module>
hls_model = convert_from_pytorch_model(
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/converters/__init__.py", line 308, in convert_from_pytorch_model
return pytorch_to_hls(config)
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/converters/pytorch_to_hls.py", line 374, in pytorch_to_hls
hls_model = ModelGraph(config, layer_list, inputs=input_layers)
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/model/graph.py", line 387, in __init__
self._make_graph(layer_list)
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/model/graph.py", line 416, in _make_graph
self.graph[name] = self.make_node(kind, name, layer, inputs, outputs)
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/model/graph.py", line 503, in make_node
node = layer_cls(self, name, attributes, inputs, outputs)
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/model/layers.py", line 117, in __init__
self.initialize()
File "/home/hls4ml-user/miniconda3/envs/hls4ml/lib/python3.10/site-packages/hls4ml/model/layers.py", line 950, in initialize
assert len(self.inputs) == 2
Optional
Possible fix
I figured out that the main reason this is happening is because the constant gets embedded in the torch.fx node representing the mul. I have been working on a fix that includes looking for constants and adding input like layers for the constants. I have gotten it to the point where the constant is represented as a layer, but haven't yet been able to get the actual constant through. I have to put it down for a few days, so I thought I'd post this to see if you think I'm on the right track. My fork is here if someone wants to have a look.
Hi! Thanks for testing this and working on a fix! I think your development goes into a promising direction. From a quick check of your fork, it seems to be like hsl4ml is still treating the new Constant layer like an input layer and is trying to find it's tensor at runtime. I'm not quite sure how best to fix it, but maybe @vloncar can advise.
A similar situation came up a while ago for RNG feature. I think the best way would be to introduce a new type of node that has no input, only output. Then we can use it as a constant input in situations like these. It would require some changes in how this is handled transparently to the rest of the nodes in the graph.
Thanks for the comment. Let me know if I can help out. For the time being I am adding additional input tensors and passing in constant values to them.
Note, we do this lots in the ONNX parsing, mostly to handle Quant nodes between the weights and where they are used. See if maybe the Constant node we have for ONNX parsing is useful here.
That new Constant node does look like it would be very useful for this use case. I'll play with that.
Yep, that seems to work. I think the Constant class that @sei-jgwohlbier implemented would work just as well. I realized that the issue in that implementation was that the new constant layer was still added to the list of input layers: https://github.com/sei-jgwohlbier/hls4ml/blob/pytorch/tensorconstant/hls4ml/converters/pytorch_to_hls.py#L292, creating the issues I saw.
I have now an implementation that seems to work and that makes use of the Constant class that was introduced with the QONNX PR at https://github.com/JanFSchulte/hls4ml/tree/constant @sei-jgwohlbier can you have a look and see if that works for you?
Oh sweet. I'll give it a shot next week. Thanks!
@JanFSchulte I'm testing this today. This update #1121 wasn't in your fork so I needed to manually add it for my case. Just FYI. I'll report back soon on whether it worked for me.
@JanFSchulte your fork worked for my test! Thanks so much!
Great, thanks for checking! PR with the fix is here: #1123
@JanFSchulte I think there might be an issue when multiple constants are present. For example, for the following class I get the same failure as before.
class test(nn.Module):
def __init__(self):
super().__init__()
self.downsample = nn.AvgPool1d(kernel_size=1, stride=2)
def forward(self, x):
d = self.downsample(x)
p = torch.mul(d,1.0)
q = torch.mul(d,2.0)
return torch.cat((p, q), dim=-1)
Ah damn, didn't think to test that. I'll have a look and update the PR accordingly.
The issue was that pytorch starts adding _n to the name of the layer for multiple instances of the same operation if n >0 , so that the matching of the operations failed. Should be fixed in the branch now.
Ok, I confirmed it works on that small test case. I'll next try my more complicated network.
Seems to be working for my resnet. Thanks!