diff --git a/dnnf/pytorch.py b/dnnf/pytorch.py index 5e29d69..196cddc 100644 --- a/dnnf/pytorch.py +++ b/dnnf/pytorch.py @@ -4,6 +4,7 @@ import torch import torch.nn.functional as F from dnnv.nn import Operation, OperationGraph, OperationVisitor, operations +from .utils import ONNX_TO_TORCH_DTYPE def convert(op_graph: OperationGraph) -> PytorchModel: @@ -235,9 +236,9 @@ def visit_Gather(self, operation: operations.Gather): def gather(operation_graph): x = torch.as_tensor(operation_graph[operation.x]) - axis = int(operation.axis) - indices = torch.as_tensor(operation_graph[operation.indices]) - result = torch.gather(x, axis, indices) + indices = [slice(None)] * x.ndim + indices[operation.axis] = operation.indices + result = x[indices] return result return gather @@ -407,12 +408,12 @@ def split(operation_graph): def visit_Sub(self, operation: operations.Sub): self.generic_visit(operation) - def add(operation_graph): + def sub(operation_graph): a = operation_graph[operation.a] b = operation_graph[operation.b] return a - b - return add + return sub def visit_Tanh(self, operation: operations.Tanh): self.generic_visit(operation) @@ -445,5 +446,87 @@ def unsqueeze(operation_graph): return unsqueeze + def visit_Upsample(self, operation: operations.Upsample): + self.generic_visit(operation) + + def upsample(operation_graph): + x = operation_graph[operation.x] + scales = operation.scales.tolist() + mode = operation.mode + result = torch.nn.Upsample(scale_factor=tuple(scales[2:]), mode=mode)(x) + return result + + return upsample + + def visit_Div(self, operation: operations.Div): + self.generic_visit(operation) + + def div(operation_graph): + a = operation_graph[operation.a] + b = operation_graph[operation.b] + result = torch.div(a, b) + return result + + return div + + def visit_Squeeze(self, operation: operations.Squeeze): + self.generic_visit(operation) + + def squeeze(operation_graph): + x = operation_graph[operation.x] + axes = operation.axes + if axes is None: + result = torch.squeeze(x) + else: + result = torch.squeeze(x, dim=axes) + return result + + return squeeze + + def visit_Expand(self, operation: operations.Expand): + self.generic_visit(operation) + + def expand(operation_graph): + x = operation_graph[operation.x] + shape = operation_graph[operation.shape] + result = x.expand(shape) + return result + + return expand + + def visit_Clip(self, operation: operations.Clip): + self.generic_visit(operation) + + def clip(operation_graph): + x = operation_graph[operation.x] + _min = operation.min + _max = operation.max + result = torch.clip(x, _min, _max) + return result + + return clip + + def visit_ReduceL2(self, operation: operations.ReduceL2): + self.generic_visit(operation) + + def reducel2(operation_graph): + x = operation_graph[operation.x] + axes = operation.axes + keepdims = operation.keepdims + result = torch.norm(x, p=2, dim=axes, keepdim=bool(keepdims)) + return result + + return reducel2 + + def visit_Cast(self, operation: operations.Cast): + self.generic_visit(operation) + + def cast(operation_graph): + x = operation_graph[operation.x] + result = x.type(ONNX_TO_TORCH_DTYPE[operation.to]) + return result + + return cast + __all__ = ["convert", "PytorchConverter"] diff --git a/dnnf/reduction.py b/dnnf/reduction.py index 8896a49..a5967c5 100644 --- a/dnnf/reduction.py +++ b/dnnf/reduction.py @@ -220,6 +220,7 @@ def suffixed_op_graph(self) -> OperationGraph: import dnnv.nn.operations as operations output_shape = self.op_graph.output_shape[0] + # axis = 0 axis = (0, 0, 1)[len(output_shape)] if len(self.op_graph.output_operations) == 1: new_output_op = self.op_graph.output_operations[0] @@ -336,22 +337,23 @@ def add_constraint(self, variables, indices, coeffs, b, is_open): self.interval_constraints[0][flat_index] = max(b / coeff, current_bound) def build(self) -> HPolyProperty: - Ab = np.vstack(self.hpoly_constraints) - A = Ab[..., :-1] - b = Ab[..., -1:] - bounds = tuple(zip(*self.interval_constraints)) - for i in range(self.num_vars): - c = np.zeros(self.num_vars) - c[i] = 1 - result = linprog(c, A, b, bounds=bounds, method="highs") - if result.success: - current_bound = self.interval_constraints[0][i] - self.interval_constraints[0][i] = max(result.x[i], current_bound) - c[i] = -1 - result = linprog(c, A, b, bounds=bounds, method="highs") - if result.success: - current_bound = self.interval_constraints[1][i] - self.interval_constraints[1][i] = min(result.x[i], current_bound) + if self.hpoly_constraints: + Ab = np.vstack(self.hpoly_constraints) + A = Ab[..., :-1] + b = Ab[..., -1:] + bounds = tuple(zip(*self.interval_constraints)) + for i in range(self.num_vars): + c = np.zeros(self.num_vars) + c[i] = 1 + result = linprog(c, A, b, bounds=bounds, method="highs") + if result.success: + current_bound = self.interval_constraints[0][i] + self.interval_constraints[0][i] = max(result.x[i], current_bound) + c[i] = -1 + result = linprog(c, A, b, bounds=bounds, method="highs") + if result.success: + current_bound = self.interval_constraints[1][i] + self.interval_constraints[1][i] = min(result.x[i], current_bound) return HPolyProperty.build( self.input_vars, self.output_vars, diff --git a/dnnf/utils.py b/dnnf/utils.py index f565814..187bc65 100644 --- a/dnnf/utils.py +++ b/dnnf/utils.py @@ -1,12 +1,28 @@ """ """ import logging +import onnx import random import sys +import torch from typing import Optional import numpy as np +ONNX_TO_TORCH_DTYPE = { + onnx.TensorProto.DOUBLE: torch.float64, + onnx.TensorProto.FLOAT16: torch.float16, + onnx.TensorProto.FLOAT: torch.float32, + onnx.TensorProto.INT8: torch.int8, + onnx.TensorProto.INT16: torch.int16, + onnx.TensorProto.INT32: torch.int32, + onnx.TensorProto.INT64: torch.int64, + onnx.TensorProto.UINT8: torch.uint8, + onnx.TensorProto.BOOL: torch.bool, + onnx.TensorProto.COMPLEX64: torch.complex64, + onnx.TensorProto.COMPLEX128: torch.complex128, +} + def set_random_seed(seed: Optional[int]) -> None: random.seed(seed)