diff --git a/src/acetone_nnet/importers/PYTORCH_importer/parse_pytorch.py b/src/acetone_nnet/importers/PYTORCH_importer/parse_pytorch.py index 2983c86..6ac2053 100644 --- a/src/acetone_nnet/importers/PYTORCH_importer/parse_pytorch.py +++ b/src/acetone_nnet/importers/PYTORCH_importer/parse_pytorch.py @@ -5,6 +5,7 @@ from acetone_nnet.generator import ( ActivationFunctions, Add, + Clip, Average, AveragePooling2D, BatchNormalization, @@ -146,11 +147,14 @@ def visit_call_function(self, node: torch.fx.Node): self.layerdic[node.name] = (this_layer, node.meta['val'].shape) self.idx+=1 elif "relu" in target_name: - self.layerdic[node.args[0].name][0].activation_function=ReLu() - self.layerdic[node.name] = self.layerdic[node.args[0].name] + self.layerdic[node.args[0].name][0].activation_function=ReLu() # set previous layer activation + self.layerdic[node.name] = self.layerdic[node.args[0].name] # register elif "clamp" in target_name: - self.layerdic[node.args[0].name][0].activation_function=ReLu() - self.layerdic[node.name] = self.layerdic[node.args[0].name] + self.layerdic[node.args[0].name][0].activation_function=Clip( # set previous layer activation + max_value=node.args[2], + min_value=node.args[1] + ) + self.layerdic[node.name] = self.layerdic[node.args[0].name] # register elif "conv2d" in target_name: stride = node.args[3][0] if len(node.args) > 3 else 1 padding = node.args[4] + node.args[4] if len(node.args) > 4 else [0,0,0,0] @@ -169,13 +173,45 @@ def visit_call_function(self, node: torch.fx.Node): input_shape=self.layerdic[node.args[0].name][1], output_shape=node.meta['val'].shape, weights=self.module.state_dict()[node.args[1].target].numpy(), - biases=self.module.state_dict()[node.args[2].target].numpy(), + biases=self.module.state_dict()[node.args[2].target].numpy() if node.args[2] is not None else np.zeros(node.meta['val'].shape[1],dtype=np.float32), activation_function=Linear(), ) self.layerdic[node.args[0].name][0].next_layer.append(this_layer) this_layer.previous_layer.append(self.layerdic[node.args[0].name][0]) self.layerdic[node.name] = (this_layer, node.meta['val'].shape) self.idx+=1 + elif "batch_norm" in target_name: + # Fusing batch normalization to previous conv2d https://pytorch-cn.com/tutorials/intermediate/fx_conv_bn_fuser.html + self.layerdic[node.name] = self.layerdic[node.args[0].name] + bn_w = self.module.state_dict()[node.args[1].target] + bn_b = self.module.state_dict()[node.args[2].target] + bn_rm = self.module.state_dict()[node.args[3].target] + bn_rv = self.module.state_dict()[node.args[4].target] + bn_eps = node.args[7] + bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps) + conv_w = torch.from_numpy(self.layerdic[node.args[0].name][0].weights) + conv_b = torch.from_numpy(self.layerdic[node.args[0].name][0].biases) + conv_w = conv_w * (bn_w * bn_var_rsqrt).reshape([-1] + [1] * (len(conv_w.shape) - 1)) + conv_b = (conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b + self.layerdic[node.args[0].name][0].weights = conv_w.numpy().astype(np.float32) + self.layerdic[node.args[0].name][0].biases = conv_b.numpy().astype(np.float32) + elif "add" in target_name: + if node.args[0].name in self.layerdic.keys(): # bypass bn2_num_batches_tracked + 1 + this_layer = Add( + original_name=node.name, + idx=self.idx, + size=node.meta['val'].numel(), + input_shapes=[self.layerdic[node.args[i].name][1] for i in range(len(node.args))], + output_shape=node.meta['val'].shape, + activation_function=Linear(), + ) + for i in range(len(node.args)): + self.layerdic[node.args[i].name][0].next_layer.append(this_layer) + this_layer.previous_layer.append(self.layerdic[node.args[i].name][0]) + self.layerdic[node.name] = (this_layer, node.meta['val'].shape) + self.idx+=1 + else: + logging.info(f"Add layer bypassed: {node.args[0].name}") elif "cat" in target_name: this_layer = Concatenate( original_name=node.name, @@ -193,6 +229,8 @@ def visit_call_function(self, node: torch.fx.Node): self.idx+=1 elif "dropout" in target_name: self.layerdic[node.name] = self.layerdic[node.args[0].name] + elif "flatten" in target_name: + self.layerdic[node.name] = self.layerdic[node.args[0].name] elif "max_pool2d" in target_name: kernel = node.args[1][0] if len(node.args) > 1 else 1 stride = node.args[2][0] if len(node.args) > 2 else 1 diff --git a/tests/models/resnet/resnet.py b/tests/models/resnet/resnet.py new file mode 100644 index 0000000..f3a0512 --- /dev/null +++ b/tests/models/resnet/resnet.py @@ -0,0 +1,286 @@ +from functools import partial +from typing import Any, Callable, Optional, Union + +import torch +import torch.nn as nn +from torch import Tensor + +class ReluN(nn.Module): + def __init__(self, n=1000) -> None: + super().__init__() + self.n = n + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.clamp(x,0,self.n) + +def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d: + """3x3 convolution with padding""" + return nn.Conv2d( + in_planes, + out_planes, + kernel_size=3, + stride=stride, + padding=dilation, + groups=groups, + bias=False, + dilation=dilation, + ) + + +def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d: + """1x1 convolution""" + return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False) + + +class BasicBlock(nn.Module): + expansion: int = 1 + + def __init__( + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + groups: int = 1, + base_width: int = 64, + dilation: int = 1, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + if groups != 1 or base_width != 64: + raise ValueError("BasicBlock only supports groups=1 and base_width=64") + if dilation > 1: + raise NotImplementedError("Dilation > 1 not supported in BasicBlock") + # Both self.conv1 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv3x3(inplanes, planes, stride) + self.bn1 = norm_layer(planes) + self.relu1 = ReluN() + self.relu2 = ReluN() + self.conv2 = conv3x3(planes, planes) + self.bn2 = norm_layer(planes) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu1(out) + + out = self.conv2(out) + out = self.bn2(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu2(out) + + return out + + +class Bottleneck(nn.Module): + # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2) + # while original implementation places the stride at the first 1x1 convolution(self.conv1) + # according to "Deep residual learning for image recognition" https://arxiv.org/abs/1512.03385. + # This variant is also known as ResNet V1.5 and improves accuracy according to + # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch. + + expansion: int = 4 + + def __init__( + self, + inplanes: int, + planes: int, + stride: int = 1, + downsample: Optional[nn.Module] = None, + groups: int = 1, + base_width: int = 64, + dilation: int = 1, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = nn.BatchNorm2d + width = int(planes * (base_width / 64.0)) * groups + # Both self.conv2 and self.downsample layers downsample the input when stride != 1 + self.conv1 = conv1x1(inplanes, width) + self.bn1 = norm_layer(width) + self.conv2 = conv3x3(width, width, stride, groups, dilation) + self.bn2 = norm_layer(width) + self.conv3 = conv1x1(width, planes * self.expansion) + self.bn3 = norm_layer(planes * self.expansion) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x: Tensor) -> Tensor: + identity = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + identity = self.downsample(x) + + out += identity + out = self.relu(out) + + return out + + +class ResNet(nn.Module): + def __init__( + self, + block: type[Union[BasicBlock, Bottleneck]], + layers: list[int], + num_classes: int = 1000, + zero_init_residual: bool = False, + groups: int = 1, + width_per_group: int = 64, + replace_stride_with_dilation: Optional[list[bool]] = None, + norm_layer: Optional[Callable[..., nn.Module]] = None, + ) -> None: + super().__init__() +# _log_api_usage_once(self) + if norm_layer is None: + norm_layer = nn.BatchNorm2d + self._norm_layer = norm_layer + + self.inplanes = 64 + self.dilation = 1 + if replace_stride_with_dilation is None: + # each element in the tuple indicates if we should replace + # the 2x2 stride with a dilated convolution instead + replace_stride_with_dilation = [False, False, False] + if len(replace_stride_with_dilation) != 3: + raise ValueError( + "replace_stride_with_dilation should be None " + f"or a 3-element tuple, got {replace_stride_with_dilation}" + ) + self.groups = groups + self.base_width = width_per_group + self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False) + self.bn1 = norm_layer(self.inplanes) + self.relu = ReluN() + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2, dilate=replace_stride_with_dilation[0]) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2, dilate=replace_stride_with_dilation[1]) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2, dilate=replace_stride_with_dilation[2]) + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): + nn.init.constant_(m.weight, 1) + nn.init.constant_(m.bias, 0) + + # Zero-initialize the last BN in each residual branch, + # so that the residual branch starts with zeros, and each residual block behaves like an identity. + # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677 + if zero_init_residual: + for m in self.modules(): + if isinstance(m, Bottleneck) and m.bn3.weight is not None: + nn.init.constant_(m.bn3.weight, 0) # type: ignore[arg-type] + elif isinstance(m, BasicBlock) and m.bn2.weight is not None: + nn.init.constant_(m.bn2.weight, 0) # type: ignore[arg-type] + + def _make_layer( + self, + block: type[Union[BasicBlock, Bottleneck]], + planes: int, + blocks: int, + stride: int = 1, + dilate: bool = False, + ) -> nn.Sequential: + norm_layer = self._norm_layer + downsample = None + previous_dilation = self.dilation + if dilate: + self.dilation *= stride + stride = 1 + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + conv1x1(self.inplanes, planes * block.expansion, stride), + norm_layer(planes * block.expansion), + ) + + layers = [] + layers.append( + block( + self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation, norm_layer + ) + ) + self.inplanes = planes * block.expansion + for _ in range(1, blocks): + layers.append( + block( + self.inplanes, + planes, + groups=self.groups, + base_width=self.base_width, + dilation=self.dilation, + norm_layer=norm_layer, + ) + ) + + return nn.Sequential(*layers) + + def _forward_impl(self, x: Tensor) -> Tensor: + # See note [TorchScript super()] + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + + return x + + def forward(self, x: Tensor) -> Tensor: + return self._forward_impl(x) + + + +def resnet18(**kwargs:Any) -> ResNet: + """ResNet-18 from `Deep Residual Learning for Image Recognition `__. + + Args: + weights (:class:`~torchvision.models.ResNet18_Weights`, optional): The + pretrained weights to use. See + :class:`~torchvision.models.ResNet18_Weights` below for + more details, and possible values. By default, no pre-trained + weights are used. + progress (bool, optional): If True, displays a progress bar of the + download to stderr. Default is True. + **kwargs: parameters passed to the ``torchvision.models.resnet.ResNet`` + base class. Please refer to the `source code + `_ + for more details about this class. + + .. autoclass:: torchvision.models.ResNet18_Weights + :members: + """ + return ResNet(BasicBlock, [2, 2, 2, 2], **kwargs) + + diff --git a/tests/models/resnet/resnet18-f37072fd.pth b/tests/models/resnet/resnet18-f37072fd.pth new file mode 100644 index 0000000..c049f92 Binary files /dev/null and b/tests/models/resnet/resnet18-f37072fd.pth differ diff --git a/tests/models/resnet/resnet18_bound.pt b/tests/models/resnet/resnet18_bound.pt new file mode 100644 index 0000000..764bf74 Binary files /dev/null and b/tests/models/resnet/resnet18_bound.pt differ diff --git a/tests/tests_inference/tests_networks/test_resnet.py b/tests/tests_inference/tests_networks/test_resnet.py index 426abcd..5da06de 100644 --- a/tests/tests_inference/tests_networks/test_resnet.py +++ b/tests/tests_inference/tests_networks/test_resnet.py @@ -22,11 +22,36 @@ from tests.common import MODELS_DIR from tests.tests_inference import acetoneTestCase - - +from torch.export import export, ExportedProgram +import torch +from tests.models.resnet.resnet import resnet18 class TestResnet(acetoneTestCase.AcetoneTestCase): """Inference test for resnet model.""" + def test_resnet_pytorch(self) -> None: + with torch.no_grad(): + #model structure + pytorch_model = resnet18() + #model weights from torchvision + pytorch_model.load_state_dict(torch.load(MODELS_DIR / "resnet" / "resnet18-f37072fd.pth")) + pytorch_model.eval() + #model Clip max values (precomputed upper bound analysis) + model_bound = torch.load(MODELS_DIR / "resnet" / "resnet18_bound.pt") + # update ReluN (Clip) upper bound + for k,v in model_bound.items(): + pytorch_model.get_submodule(k).n = v + data = torch.rand(1,3,224,224, requires_grad=False, dtype=torch.float32) + program : ExportedProgram = export(pytorch_model,(data,)) + acetone_result,python_result = acetoneTestCase.run_acetone_for_test( + self.tmpdir_name, + program, + data.numpy(), + bin_dataset=True + ) + self.assertListAlmostEqual(acetone_result, python_result) + self.assertListAlmostEqual(pytorch_model(data).detach().numpy()[0], python_result) + + def test_resnet_onnx(self) -> None: """Test resnet model, compare between keras et C code.""" model_path = MODELS_DIR / "resnet" / "resnet18-v2-7.onnx" diff --git a/tests/tests_inference/tests_networks/test_squeezenet.py b/tests/tests_inference/tests_networks/test_squeezenet.py index c77964d..00152b0 100644 --- a/tests/tests_inference/tests_networks/test_squeezenet.py +++ b/tests/tests_inference/tests_networks/test_squeezenet.py @@ -89,8 +89,9 @@ class TestSqueezenet(acetoneTestCase.AcetoneTestCase): def test_squeezenet_pytorch(self) -> None: with torch.no_grad(): - data = torch.rand(1,3,224,224, requires_grad=False) + data = torch.rand(1,3,224,224, requires_grad=False, dtype=torch.float32) pytorch_model = SqueezeNetv11(relun=1000) + pytorch_model.eval() program : ExportedProgram = export(pytorch_model,(data,)) acetone_result,python_result = acetoneTestCase.run_acetone_for_test( self.tmpdir_name,