How to use the nemo.core.AxisType function in NEMO

To help you get started, we’ve selected a few NEMO examples, based on popular ways it is used in public projects.

Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.

github NVIDIA / NeMo / examples / image / mnist_lenet5.py View on Github external
def create_ports():
        input_ports = {
            "images": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 1),
                                  2: AxisType(HeightTag, 32),
                                  3: AxisType(WidthTag, 32)
                                  })
        }
        output_ports = {
            "predictions": NeuralType({0: AxisType(BatchTag),
                                       1: AxisType(LogProbabilityTag)
                                       })
            # no PredictionTag!?

        }
        return input_ports, output_ports
github NVIDIA / NeMo / collections / nemo_simple_gan / nemo_simple_gan / gan.py View on Github external
def create_ports():
        input_ports = {
            "latents": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 64),
                                  2: AxisType(HeightTag, 4),
                                  3: AxisType(WidthTag, 4)})
        }
        output_ports = {
            "image": NeuralType({0: AxisType(BatchTag),
                                 1: AxisType(ChannelTag),
                                 2: AxisType(HeightTag, 28),
                                 3: AxisType(WidthTag, 28)})
        }
        return input_ports, output_ports
github NVIDIA / NeMo / collections / nemo_simple_gan / nemo_simple_gan / gan.py View on Github external
def create_ports(input_size=(32, 32)):
        input_ports = {}
        output_ports = {
            "latent": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 64),
                                  2: AxisType(HeightTag, 4),
                                  3: AxisType(WidthTag, 4)}),
            "image": NeuralType({0: AxisType(BatchTag),
                                 1: AxisType(ChannelTag),
                                 2: AxisType(HeightTag, input_size[1]),
                                 3: AxisType(WidthTag, input_size[0])}),
            "label": NeuralType({0: AxisType(BatchTag)})
        }
        return input_ports, output_ports
github NVIDIA / NeMo / collections / nemo_simple_gan / nemo_simple_gan / gan.py View on Github external
def create_ports():
        input_ports = {
            "latents": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 64),
                                  2: AxisType(HeightTag, 4),
                                  3: AxisType(WidthTag, 4)})
        }
        output_ports = {
            "image": NeuralType({0: AxisType(BatchTag),
                                 1: AxisType(ChannelTag),
                                 2: AxisType(HeightTag, 28),
                                 3: AxisType(WidthTag, 28)})
        }
        return input_ports, output_ports
github NVIDIA / NeMo / examples / image / mnist_lenet5.py View on Github external
def create_ports():
        input_ports = {
            "images": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 1),
                                  2: AxisType(HeightTag, 32),
                                  3: AxisType(WidthTag, 32)
                                  })
        }
        output_ports = {
            "predictions": NeuralType({0: AxisType(BatchTag),
                                       1: AxisType(LogProbabilityTag)
                                       })
            # no PredictionTag!?

        }
        return input_ports, output_ports
github NVIDIA / NeMo / examples / image / mnist_lenet5.py View on Github external
def create_ports(input_size=(32, 32)):
        input_ports = {}
        output_ports = {
            "images": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 1),
                                  2: AxisType(HeightTag, input_size[1]),
                                  3: AxisType(WidthTag, input_size[0])}),
            "targets": NeuralType({0: AxisType(BatchTag)})
        }
        return input_ports, output_ports
github NVIDIA / NeMo / examples / image / mnist_lenet5.py View on Github external
def create_ports():
        input_ports = {
            "predictions": NeuralType(
                {0: AxisType(BatchTag),
                 1: AxisType(LogProbabilityTag)}),
            "targets": NeuralType({0: AxisType(BatchTag)}),
        }
        output_ports = {"loss": NeuralType(None)}
        return input_ports, output_ports
github NVIDIA / NeMo / collections / nemo_simple_gan / nemo_simple_gan / gan.py View on Github external
def create_ports():
        input_ports = {
            "image1": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag),
                                  2: AxisType(HeightTag, 28),
                                  3: AxisType(WidthTag, 28)}),
            "image2": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag),
                                  2: AxisType(HeightTag, 28),
                                  3: AxisType(WidthTag, 28)})
        }

        output_ports = {
            "interpolated_image": NeuralType({0: AxisType(BatchTag),
                                              1: AxisType(ChannelTag),
                                              2: AxisType(HeightTag, 28),
                                              3: AxisType(WidthTag, 28)})
        }
        return input_ports, output_ports
github NVIDIA / NeMo / collections / nemo_simple_gan / nemo_simple_gan / gan.py View on Github external
def create_ports(input_size=(32, 32)):
        input_ports = {}
        output_ports = {
            "latent": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag, 64),
                                  2: AxisType(HeightTag, 4),
                                  3: AxisType(WidthTag, 4)}),
            "image": NeuralType({0: AxisType(BatchTag),
                                 1: AxisType(ChannelTag),
                                 2: AxisType(HeightTag, input_size[1]),
                                 3: AxisType(WidthTag, input_size[0])}),
            "label": NeuralType({0: AxisType(BatchTag)})
        }
        return input_ports, output_ports
github NVIDIA / NeMo / collections / nemo_simple_gan / nemo_simple_gan / gan.py View on Github external
def create_ports():
        input_ports = {
            "image1": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag),
                                  2: AxisType(HeightTag, 28),
                                  3: AxisType(WidthTag, 28)}),
            "image2": NeuralType({0: AxisType(BatchTag),
                                  1: AxisType(ChannelTag),
                                  2: AxisType(HeightTag, 28),
                                  3: AxisType(WidthTag, 28)})
        }

        output_ports = {
            "interpolated_image": NeuralType({0: AxisType(BatchTag),
                                              1: AxisType(ChannelTag),
                                              2: AxisType(HeightTag, 28),
                                              3: AxisType(WidthTag, 28)})
        }
        return input_ports, output_ports