<commit_before>#!/usr/bin/python

# Define a stack to host a Consul cluster
#
# Template Parameters (provided at templat creation time):
# - name
#   Name of the stack, of course. Required.
# - description
#   Description of the stack. Please provide one.
# - vpc_id
#   ID of the VPC in which this stack is built
# - vpc_cidr
#   CIDR block of the vpc
# - subnet_ids
#   List of subnets in which the consul cluster should be built.
# - cluster_size
#   How many instances should participate in the cluster? Should be at least 3.
# - instance_type
#   Instance type of the instances. Defaults to t2.micro
#
# Stack Parameters (provided to the template at stack create/update time):
#
# - ConsulKey
#   Name of the key pair to use to connect to Consul cluster server instances
#
# Stack Outputs:
#

import troposphere.ec2 as ec2
import troposphere.iam as iam
import troposphere.autoscaling as asg
import troposphere.cloudformation as cf
import troposphere.logs as logs
import troposphere as tp

from ..template import asgtag, TemplateBuilder, AMI_REGION_MAP_NAME, REF_STACK_NAME
from ..network import net

class ConsulTemplate(TemplateBuilder):

    BUILD_PARM_NAMES = ['region', 'bucket', 'vpc_id', 'vpc_cidr', 'subnet_ids', 'cluster_size', 'instance_type']
    CONSUL_KEY_PARAM_NAME = 'ConsulKey'

    def __init__(self, name,
                 region,
                 bucket,
                 vpc_id,
                 vpc_cidr,
                 subnet_ids,
                 description='[REPLACEME]',
                 cluster_size=3,
                 instance_type='t2.micro'):
        super(ConsulTemplate, self).__init__(name, description, ConsulTemplate.BUILD_PARM_NAMES)

        self.region = region
        self.bucket = bucket
        self.vpc_id = vpc_id
        self.vpc_cidr = vpc_cidr
        self.subnet_ids = subnet_ids
        self.cluster_size = cluster_size
        self.instance_type = instance_type

    def build_template(self):
        super(ConsulTemplate, self).build_template()

        self.create_parameters()

        sg = self.create_sg()
        iam_profile = self.create_iam_profile()
        self.create_logstreams()

        self.enis = [
            self.create_eni(i, sg, self.subnet_ids[i % len(self.subnet_ids)]) for i in range(0, self.cluster_size)
        ]
        self.asgs = [
            self.create_asg(i, sg, iam_profile, self.subnet_ids[i % len(self.subnet_ids)], self.enis[i]) for i in range(self.cluster_size)
        ]

    def create_parameters(self):
        self.add_parameter(tp.Parameter(self.CONSUL_KEY_PARAM_NAME, Type='String'))

    def create_eni(self, index, security_group, subnet_id):
        eni = ec2.NetworkInterface('Consul{}ENI'.format(index),
                                   Description='ENI for Consul cluster member {}'.format(index),
                                   GroupSet=[tp.Ref(security_group)],
                                   SourceDestCheck=True,
                                   SubnetId=subnet_id,
                                   Tags=self.default_tags)
        self.add_resource(eni)
        self.add_output(tp.Output(eni.name, Value=tp.Ref(eni)))
        return eni

    def create_sg(self):
        rules = [net.sg_rule(self.vpc_cidr, net.SSH, net.TCP)] + [
            net.sg_rule(self.vpc_cidr, ports, protocol)
            for protocol in [net.TCP, net.UDP] for ports in [53, (8300, 8302), 8400, 8500, 8600]
        ]
        sg = ec2.SecurityGroup('ConsulSecurityGroup',
                               GroupDescription='Consul Instance Security Group',
                               SecurityGroupIngress=rules,
                               VpcId=self.vpc_id,
                               Tags=self.default_tags)
        self.add_resource(sg)
        return sg

    def create_asg(self, index, security_group, iam_profile, subnet_id, eni):
        lc_name = 'Consul{}LC'.format(index)
        lc = asg.LaunchConfiguration(lc_name,
                                     ImageId=tp.FindInMap(AMI_REGION_MAP_NAME, self.region, 'GENERAL'),
                                     InstanceType=self.instance_type,
                                     SecurityGroups=[tp.Ref(security_group)],
                                     KeyName=tp.Ref(self.CONSUL_KEY_PARAM_NAME),
                                     IamInstanceProfile=tp.Ref(iam_profile),
                                     InstanceMonitoring=False,
                                     AssociatePublicIpAddress=False,
                                     UserData=self._create_consul_userdata(eni, lc_name),
                                     Metadata=self._create_consul_metadata(eni, index))
        group = asg.AutoScalingGroup('Consul{}ASG'.format(index),
                                     MinSize=1, MaxSize=1,
                                     LaunchConfigurationName=tp.Ref(lc),
                                     VPCZoneIdentifier=[subnet_id],
                                     Tags=asgtag(self._rename('{} Server-' + str(index))))
        self.add_resources(lc, group)
        self.add_output(tp.Output(group.name, Value=tp.Ref(group)))
        return group

    def create_logstreams(self):
        self.log_group = logs.LogGroup('ConsulLogGroup',
                                       RetentionInDays=3)
        self.add_resources(self.log_group)

    def create_iam_profile(self):
        role = iam.Role('ConsulInstanceRole',
                        AssumeRolePolicyDocument={
                            'Statement': [{
                                'Effect': 'Allow',
                                'Principal': {'Service': ['ec2.amazonaws.com']},
                                'Action': ['sts:AssumeRole']
                            }]
                        },
                        Policies=[
                            iam.Policy(
                                PolicyName='ConsulInstance',
                                PolicyDocument={
                                    'Statement': [{
                                        'Effect': 'Allow',
                                        'Resource': ['*'],
                                        'Action': ['ec2:Attach*', 'ec2:Describe*', 's3:*']
                                    }, {
                                        'Effect': 'Allow',
                                        'Resource': 'arn:aws:s3:::{}/*'.format(self.bucket),
                                        'Action': ['s3:Get*']
                                    }, {
                                        'Effect': 'Allow',
                                        'Resource': 'arn:aws:logs:{}:*:*'.format(self.region),
                                        'Action': [
                                            'logs:CreateLogGroup',
                                            'logs:CreateLogStream',
                                            'logs:PutLogEvents',
                                            'logs:DescribeLogStreams'
                                        ]
                                    }]
                                }
                            )
                        ])
        profile = iam.InstanceProfile('ConsulInstanceProfile',
                                      Path='/',
                                      Roles=[tp.Ref(role)])
        self.add_resources(role, profile)
        return profile

    def _create_consul_userdata(self, eni, resource_name):
        startup = [
            '#!/bin/bash\n',
            'yum update -y && yum install -y yum-cron && chkconfig yum-cron on\n',
            'REGION=', self.region, '\n',
            'INS_ID=$(curl http://169.254.169.254/latest/meta-data/instance-id)\n',
            'ENI_ID=', tp.Ref(eni), '\n',
            # TODO: do we need a sleep/check loop here? It hasn't failed so far
            'aws ec2 attach-network-interface --instance-id $INS_ID --device-index 1 --network-interface-id $ENI_ID --region $REGION\n',
            '/opt/aws/bin/cfn-init -v ',
            '  --stack ', REF_STACK_NAME,
            '  --resource ', resource_name,
            '  --configsets install',
            '  --region $REGION\n',
        ]
        return tp.Base64(tp.Join('', startup))  # TODO: There has GOT to be a better way to do userdata.

    def _create_consul_metadata(self, eni, index):
        consul_dir = '/opt/consul'
        agent_dir = '{}/agent'.format(consul_dir)
        config_dir = '{}/config'.format(consul_dir)
        config_file = '{}/config.json'.format(config_dir)
        config_py = '{}/config.py'.format(consul_dir)
        data_dir = '{}/data'.format(consul_dir)
        source_prefix = 'http://{}.s3.amazonaws.com/scaffold/consul'.format(self.bucket)
        return cf.Metadata(
            cf.Init(
                cf.InitConfigSets(install=['install']),
                install=cf.InitConfig(
                    packages={
                        'python': {
                            'boto3': []
                        }
                    },
                    # groups={}, # do we need a consul group?
                    # users={}, # do we need a consul user?
                    sources={
                        agent_dir: 'https://releases.hashicorp.com/consul/0.6.3/consul_0.6.3_linux_amd64.zip',
                    },
                    files={
                        # See https://www.consul.io/docs/agent/options.html#configuration_files
                        config_file: {
                            'content': {
                                'datacenter': self.region,
                                'data_dir': data_dir,
                                'log_level': 'INFO',
                                'server': True,
                                'bootstrap_expect': self.cluster_size,
                                'bind_addr': tp.Ref(eni),  # resolved later
                                'retry_join': [tp.Ref(e) for e in self.enis if e != eni]  # resolved later
                            },
                            'mode': '000755',
                            'owner': 'root',  # could be consul user?
                            'group': 'root'   # could be consul group?
                        },
                        config_py: {
                            'source': '{}/config.py'.format(source_prefix),
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        },
                        '/etc/init.d/consul': {
                            'source': '{}/consul.service'.format(source_prefix),
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        },
                        '/opt/cw-logs/awslogs-agent-setup.py': {
                            'source': 'https://s3.amazonaws.com/aws-cloudwatch/downloads/latest/awslogs-agent-setup.py',
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        },
                        '/opt/cw-logs/cwlogs.cfg': {
                            'content': tp.Join('', [
                                '[general]\n',
                                'state_file = /var/awslogs/state/agent-state\n',
                                '\n',
                                '[/var/log/consul]\n',
                                'file = /var/log/consul\n',
                                'datetime_format = %b %d %H:%M:%S\n',
                                'log_group_name = ', tp.Ref(self.log_group), '\n',
                                'log_stream_name = consul_server_{}\n'.format(index),
                                ]),
                            'source': '{}/cwlogs.cfg'.format(source_prefix),
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        }
                    },
                    commands={
                        '10_dirs': {
                            'command': 'mkdir -p {}'.format(data_dir)  # also -m 755?
                        },
                        '20_mode': {
                            'command': 'chmod 755 {}/consul'.format(agent_dir)
                        },
                        '30_config': {
                            'command': 'cp {0} {0}.orig && python {1} {0}'.format(config_file, config_py)  # blech.
                        },
                        '40_wait': {
                            'command': 'while [ `ifconfig | grep "inet addr" | wc -l` -lt 3 ]; do echo "waiting for ip addr" > /opt/consul/wait && sleep 2; done'
                        },
                        '50_chkconfig': {
                            'command': 'chkconfig --add consul'
                        },
                        '60_cwlogs': {
                            'command': 'python /opt/cw-logs/awslogs-agent-setup.py -n -r {} -c /opt/cw-logs/cwlogs.cfg'.format(self.region)
                        }
                    },
                    services={
                        'sysvinit': {
                            'consul': {
                                'enabled': 'true',
                                'ensureRunning': 'true',
                                'files': [config_file],
                                'sources': [agent_dir],
                                'commands': 'config'
                            }
                        },
                        # dnsmasq?
                    }
                )
            )
        )

if __name__ == '__main__':
    import sys
    name = sys.argv[1] if len(sys.argv) > 1 else 'Test'
    region = sys.argv[2] if len(sys.argv) > 2 else 'us-west-2'
    bucket = sys.argv[3] if len(sys.argv) > 3 else 'thousandleaves-us-west-2-laurel-deploy'
    vpc_id = sys.argv[4] if len(sys.argv) > 4 else 'vpc-deadbeef'
    vpc_cidr = sys.argv[5] if len(sys.argv) > 5 else '10.0.0.0/16'
    subnet_ids = sys.argv[6:] if len(sys.argv) > 6 else ['subnet-deadbeef', 'subnet-cab4abba']

    template = ConsulTemplate(name, region, bucket, vpc_id, vpc_cidr, subnet_ids)
    template.build_template()
    print template.to_json()
<commit_msg>add some missing stack parameters and ouputs<commit_after>#!/usr/bin/python

# Define a stack to host a Consul cluster
#
# Template Parameters (provided at templat creation time):
# - name
#   Name of the stack, of course. Required.
# - description
#   Description of the stack. Please provide one.
# - region
#   The AWS region in which this template will be executed
# - bucket
#   The S3 bucket to be used for uploading server configuration files
# - vpc_id
#   ID of the VPC in which this stack is built
# - vpc_cidr
#   CIDR block of the vpc
# - subnet_ids
#   List of subnets in which the consul cluster should be built.
# - cluster_size
#   How many instances should participate in the cluster? Should be at least 3.
# - instance_type
#   Instance type of the instances. Defaults to t2.micro
#
# Stack Parameters (provided to the template at stack create/update time):
#
# - ConsulKey
#   Name of the key pair to use to connect to Consul cluster server instances
#
# Stack Outputs:
#
# - Consul{N}ASG
#   ID of the autoscaling group controlling the Nth cluster member
# - Consul{N}ENI
#   ID of the elastic network interface attached to the Nth cluster member

import troposphere.ec2 as ec2
import troposphere.iam as iam
import troposphere.autoscaling as asg
import troposphere.cloudformation as cf
import troposphere.logs as logs
import troposphere as tp

from ..template import asgtag, TemplateBuilder, AMI_REGION_MAP_NAME, REF_STACK_NAME
from ..network import net

class ConsulTemplate(TemplateBuilder):

    BUILD_PARM_NAMES = ['region', 'bucket', 'vpc_id', 'vpc_cidr', 'subnet_ids', 'cluster_size', 'instance_type']
    CONSUL_KEY_PARAM_NAME = 'ConsulKey'

    def __init__(self, name,
                 region,
                 bucket,
                 vpc_id,
                 vpc_cidr,
                 subnet_ids,
                 description='[REPLACEME]',
                 cluster_size=3,
                 instance_type='t2.micro'):
        super(ConsulTemplate, self).__init__(name, description, ConsulTemplate.BUILD_PARM_NAMES)

        self.region = region
        self.bucket = bucket
        self.vpc_id = vpc_id
        self.vpc_cidr = vpc_cidr
        self.subnet_ids = subnet_ids
        self.cluster_size = cluster_size
        self.instance_type = instance_type

    def build_template(self):
        super(ConsulTemplate, self).build_template()

        self.create_parameters()

        sg = self.create_sg()
        iam_profile = self.create_iam_profile()
        self.create_logstreams()

        self.enis = [
            self.create_eni(i, sg, self.subnet_ids[i % len(self.subnet_ids)]) for i in range(0, self.cluster_size)
        ]
        self.asgs = [
            self.create_asg(i, sg, iam_profile, self.subnet_ids[i % len(self.subnet_ids)], self.enis[i]) for i in range(self.cluster_size)
        ]

    def create_parameters(self):
        self.add_parameter(tp.Parameter(self.CONSUL_KEY_PARAM_NAME, Type='String'))

    def create_eni(self, index, security_group, subnet_id):
        eni = ec2.NetworkInterface('Consul{}ENI'.format(index),
                                   Description='ENI for Consul cluster member {}'.format(index),
                                   GroupSet=[tp.Ref(security_group)],
                                   SourceDestCheck=True,
                                   SubnetId=subnet_id,
                                   Tags=self.default_tags)
        self.add_resource(eni)
        self.add_output(tp.Output(eni.name, Value=tp.Ref(eni)))
        return eni

    def create_sg(self):
        rules = [net.sg_rule(self.vpc_cidr, net.SSH, net.TCP)] + [
            net.sg_rule(self.vpc_cidr, ports, protocol)
            for protocol in [net.TCP, net.UDP] for ports in [53, (8300, 8302), 8400, 8500, 8600]
        ]
        sg = ec2.SecurityGroup('ConsulSecurityGroup',
                               GroupDescription='Consul Instance Security Group',
                               SecurityGroupIngress=rules,
                               VpcId=self.vpc_id,
                               Tags=self.default_tags)
        self.add_resource(sg)
        return sg

    def create_asg(self, index, security_group, iam_profile, subnet_id, eni):
        lc_name = 'Consul{}LC'.format(index)
        lc = asg.LaunchConfiguration(lc_name,
                                     ImageId=tp.FindInMap(AMI_REGION_MAP_NAME, self.region, 'GENERAL'),
                                     InstanceType=self.instance_type,
                                     SecurityGroups=[tp.Ref(security_group)],
                                     KeyName=tp.Ref(self.CONSUL_KEY_PARAM_NAME),
                                     IamInstanceProfile=tp.Ref(iam_profile),
                                     InstanceMonitoring=False,
                                     AssociatePublicIpAddress=False,
                                     UserData=self._create_consul_userdata(eni, lc_name),
                                     Metadata=self._create_consul_metadata(eni, index))
        group = asg.AutoScalingGroup('Consul{}ASG'.format(index),
                                     MinSize=1, MaxSize=1,
                                     LaunchConfigurationName=tp.Ref(lc),
                                     VPCZoneIdentifier=[subnet_id],
                                     Tags=asgtag(self._rename('{} Server-' + str(index))))
        self.add_resources(lc, group)
        self.add_output(tp.Output(group.name, Value=tp.Ref(group)))
        return group

    def create_logstreams(self):
        self.log_group = logs.LogGroup('ConsulLogGroup',
                                       RetentionInDays=3)
        self.add_resources(self.log_group)

    def create_iam_profile(self):
        role = iam.Role('ConsulInstanceRole',
                        AssumeRolePolicyDocument={
                            'Statement': [{
                                'Effect': 'Allow',
                                'Principal': {'Service': ['ec2.amazonaws.com']},
                                'Action': ['sts:AssumeRole']
                            }]
                        },
                        Policies=[
                            iam.Policy(
                                PolicyName='ConsulInstance',
                                PolicyDocument={
                                    'Statement': [{
                                        'Effect': 'Allow',
                                        'Resource': ['*'],
                                        'Action': ['ec2:Attach*', 'ec2:Describe*', 's3:*']
                                    }, {
                                        'Effect': 'Allow',
                                        'Resource': 'arn:aws:s3:::{}/*'.format(self.bucket),
                                        'Action': ['s3:Get*']
                                    }, {
                                        'Effect': 'Allow',
                                        'Resource': 'arn:aws:logs:{}:*:*'.format(self.region),
                                        'Action': [
                                            'logs:CreateLogGroup',
                                            'logs:CreateLogStream',
                                            'logs:PutLogEvents',
                                            'logs:DescribeLogStreams'
                                        ]
                                    }]
                                }
                            )
                        ])
        profile = iam.InstanceProfile('ConsulInstanceProfile',
                                      Path='/',
                                      Roles=[tp.Ref(role)])
        self.add_resources(role, profile)
        return profile

    def _create_consul_userdata(self, eni, resource_name):
        startup = [
            '#!/bin/bash\n',
            'yum update -y && yum install -y yum-cron && chkconfig yum-cron on\n',
            'REGION=', self.region, '\n',
            'INS_ID=$(curl http://169.254.169.254/latest/meta-data/instance-id)\n',
            'ENI_ID=', tp.Ref(eni), '\n',
            # TODO: do we need a sleep/check loop here? It hasn't failed so far
            'aws ec2 attach-network-interface --instance-id $INS_ID --device-index 1 --network-interface-id $ENI_ID --region $REGION\n',
            '/opt/aws/bin/cfn-init -v ',
            '  --stack ', REF_STACK_NAME,
            '  --resource ', resource_name,
            '  --configsets install',
            '  --region $REGION\n',
        ]
        return tp.Base64(tp.Join('', startup))  # TODO: There has GOT to be a better way to do userdata.

    def _create_consul_metadata(self, eni, index):
        consul_dir = '/opt/consul'
        agent_dir = '{}/agent'.format(consul_dir)
        config_dir = '{}/config'.format(consul_dir)
        config_file = '{}/config.json'.format(config_dir)
        config_py = '{}/config.py'.format(consul_dir)
        data_dir = '{}/data'.format(consul_dir)
        source_prefix = 'http://{}.s3.amazonaws.com/scaffold/consul'.format(self.bucket)
        return cf.Metadata(
            cf.Init(
                cf.InitConfigSets(install=['install']),
                install=cf.InitConfig(
                    packages={
                        'python': {
                            'boto3': []
                        }
                    },
                    # groups={}, # do we need a consul group?
                    # users={}, # do we need a consul user?
                    sources={
                        agent_dir: 'https://releases.hashicorp.com/consul/0.6.3/consul_0.6.3_linux_amd64.zip',
                    },
                    files={
                        # See https://www.consul.io/docs/agent/options.html#configuration_files
                        config_file: {
                            'content': {
                                'datacenter': self.region,
                                'data_dir': data_dir,
                                'log_level': 'INFO',
                                'server': True,
                                'bootstrap_expect': self.cluster_size,
                                'bind_addr': tp.Ref(eni),  # resolved later
                                'retry_join': [tp.Ref(e) for e in self.enis if e != eni]  # resolved later
                            },
                            'mode': '000755',
                            'owner': 'root',  # could be consul user?
                            'group': 'root'   # could be consul group?
                        },
                        config_py: {
                            'source': '{}/config.py'.format(source_prefix),
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        },
                        '/etc/init.d/consul': {
                            'source': '{}/consul.service'.format(source_prefix),
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        },
                        '/opt/cw-logs/awslogs-agent-setup.py': {
                            'source': 'https://s3.amazonaws.com/aws-cloudwatch/downloads/latest/awslogs-agent-setup.py',
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        },
                        '/opt/cw-logs/cwlogs.cfg': {
                            'content': tp.Join('', [
                                '[general]\n',
                                'state_file = /var/awslogs/state/agent-state\n',
                                '\n',
                                '[/var/log/consul]\n',
                                'file = /var/log/consul\n',
                                'datetime_format = %b %d %H:%M:%S\n',
                                'log_group_name = ', tp.Ref(self.log_group), '\n',
                                'log_stream_name = consul_server_{}\n'.format(index),
                                ]),
                            'source': '{}/cwlogs.cfg'.format(source_prefix),
                            'mode': '000755',
                            'owner': 'root',
                            'group': 'root'
                        }
                    },
                    commands={
                        '10_dirs': {
                            'command': 'mkdir -p {}'.format(data_dir)  # also -m 755?
                        },
                        '20_mode': {
                            'command': 'chmod 755 {}/consul'.format(agent_dir)
                        },
                        '30_config': {
                            'command': 'cp {0} {0}.orig && python {1} {0}'.format(config_file, config_py)  # blech.
                        },
                        '40_wait': {
                            'command': 'while [ `ifconfig | grep "inet addr" | wc -l` -lt 3 ]; do echo "waiting for ip addr" > /opt/consul/wait && sleep 2; done'
                        },
                        '50_chkconfig': {
                            'command': 'chkconfig --add consul'
                        },
                        '60_cwlogs': {
                            'command': 'python /opt/cw-logs/awslogs-agent-setup.py -n -r {} -c /opt/cw-logs/cwlogs.cfg'.format(self.region)
                        }
                    },
                    services={
                        'sysvinit': {
                            'consul': {
                                'enabled': 'true',
                                'ensureRunning': 'true',
                                'files': [config_file],
                                'sources': [agent_dir],
                                'commands': 'config'
                            }
                        },
                        # dnsmasq?
                    }
                )
            )
        )

if __name__ == '__main__':
    import sys
    name = sys.argv[1] if len(sys.argv) > 1 else 'Test'
    region = sys.argv[2] if len(sys.argv) > 2 else 'us-west-2'
    bucket = sys.argv[3] if len(sys.argv) > 3 else 'thousandleaves-us-west-2-laurel-deploy'
    vpc_id = sys.argv[4] if len(sys.argv) > 4 else 'vpc-deadbeef'
    vpc_cidr = sys.argv[5] if len(sys.argv) > 5 else '10.0.0.0/16'
    subnet_ids = sys.argv[6:] if len(sys.argv) > 6 else ['subnet-deadbeef', 'subnet-cab4abba']

    template = ConsulTemplate(name, region, bucket, vpc_id, vpc_cidr, subnet_ids)
    template.build_template()
    print template.to_json()
<|endoftext|>


<commit_before>import inviwopy
import numpy as np
import h5py
from envisionpy.utils.exceptions import *
from envisionpy.utils.atomData import atomic_radii, element_names, element_colors
from .baseNetworks.Decoration import Decoration
# TODO add volume merger and multi-raycaster

class AtomPositions(Decoration):
    '''
    Manages a subnetwork for atom position rendering.
    '''
    def __init__(self, inviwoApp, hdf5_path, hdf5_output, xpos=0, ypos=0):
        Decoration.__init__(self, inviwoApp)
        self.atom_radii = []
        self.atom_names = []
        self.nAtomTypes = 0
        self.setup_network(hdf5_path, hdf5_output, xpos, ypos)

        self.set_atom_radius(0.5)
        self.hide()

    @staticmethod
    def valid_hdf5(hdf5_file):
        return hdf5_file.get("UnitCell") != None

    def valid_visualisations(self):
        return ['charge', 'elf', 'parchg']

    def connect_decoration(self, other, vis_type):
        # Add a decoration by connecting data ports and linking properties.
        if vis_type not in self.valid_visualisations():
            raise EnvisionError('Invalid visualisation type ['+vis_type+'].')

        # Link needed properties between networks.
        if vis_type == 'charge' or vis_type == 'elf' or vis_type == 'parchg':
            self.network.addLink(self.camera_prop, other.camera_prop)
            self.network.addLink(other.camera_prop, self.camera_prop)

        other.connect_decoration_ports(self.decoration_outport)

    def disconnect_decoration(self, other, vis_type):
        if vis_type == 'charge' or vis_type == 'elf' or vis_type == 'parchg':
            self.network.removeLink(self.camera_prop, other.camera_prop)
        other.disconnect_decorations_port(self.decoration_outport)

    def show(self):
        self.get_processor('UnitcellCanvas').widget.show()

    def hide(self):
        self.get_processor('UnitcellCanvas').widget.hide()

    # def get_ui_data(self):
    #     rc = self.get_processor('Raycaster')
    #     vCanvas = self.get_processor('VolumeCanvas')
    #     sCanvas = self.get_processor('SliceCanvas')
    #     selection = self.get_processor('Hdf5Selection')
    #     volumeSlice = self.get_processor('VolumeSlice')
# ------------------------------------------
# ------- Property control functions -------

    def set_atom_radius(self, radius, index=None):
        structureMesh = self.get_processor('UnitcellMesh')
        if structureMesh.fullMesh.value:
            if index != None:
                self.atom_radii[index] = radius
                structureMesh.getPropertyByIdentifier("radius" + str(index)).value = radius
            else:
                for i in range(self.nAtomTypes):
                    self.atom_radii[i] = radius
                    self.set_atom_radius(radius, i)
        else:
            sphereRenderer = self.get_processor('UnitcellRenderer')
            sphereRenderer.sphereProperties.defaultRadius.value = radius
            for i in range(self.nAtomTypes):
                self.atom_radii[i] = radius

    def hide_atoms(self):
        return self.set_atom_radius(0)

    def get_atom_names(self):
        return self.atom_names

    def get_atom_name(self, index):
        return self.atom_names[index]

    def get_atom_radii(self):
        return self.atom_names

    def toggle_full_mesh(self, enable):
        structMesh = self.get_processor('UnitcellMesh')
        structMesh.fullMesh.value = enable

    def set_canvas_position(self, x, y):
    # Updates the position of the canvas
    # Upper left corner will be at coordinate (x, y)
        unitcellCanvas = self.get_processor('UnitcellCanvas')
        unitcellCanvas.position.value = inviwopy.glm.ivec2(x, y)

# ------------------------------------------
# ------- Network building functions -------

    def setup_network(self, hdf5_path, hdf5_output, xpos, ypos):
        strucMesh = self.add_processor('envision.StructureMesh', 'UnitcellMesh', xpos, ypos+3)
        meshRenderer = self.add_processor('org.inviwo.SphereRenderer', 'UnitcellRenderer', xpos, ypos+6)
        background = self.add_processor('org.inviwo.Background', 'AtomBackground', xpos, ypos+9)
        canvas = self.add_processor('org.inviwo.CanvasGL', 'UnitcellCanvas', xpos, ypos+12)

        canvas.inputSize.dimensions.value = inviwopy.glm.size2_t(500, 500)

        self.network.addConnection(strucMesh.getOutport('mesh'), meshRenderer.getInport('geometry'))
        self.network.addConnection(meshRenderer.getOutport('image'), background.getInport('inport'))
        self.network.addConnection(background.getOutport('outport'), canvas.getInport('inport'))

        with h5py.File(hdf5_path, "r") as h5:
            # Set basis matrix and scaling
            basis_matrix = np.array(h5["/basis"], dtype='d')
            strucMesh.basis.minValue = inviwopy.glm.mat3(
                -1000, -1000, -1000,
                -1000, -1000, -1000,
                -1000, -1000, -1000)
            strucMesh.basis.maxValue = inviwopy.glm.mat3(
                1000,1000,1000,
                1000,1000,1000,
                1000,1000,1000)
            strucMesh.basis.value = inviwopy.glm.mat3(
                basis_matrix[0,0], basis_matrix[0,1], basis_matrix[0,2],
                basis_matrix[1,0], basis_matrix[1,1], basis_matrix[1,2],
                basis_matrix[2,0], basis_matrix[2,1], basis_matrix[2,2])
            strucMesh.scalingFactor.maxValue = h5['/scaling_factor'][()]
            strucMesh.scalingFactor.value = h5['/scaling_factor'][()]

            # Setup animation?
            strucMesh.timestep.value = 0
            strucMesh.timestep.minValue = 0
            base_group = "/UnitCell"

            strucMesh.timestep.maxValue = 0
            for i,key in enumerate(list(h5[base_group + "/Atoms"].keys())):

                element = h5[base_group + "/Atoms/"+key].attrs['element']

                name = element_names.get(element, 'Unknown')
                self.atom_names.append(name)
                color = element_colors.get(element, (0.5, 0.5, 0.5, 1.0))
                radius = atomic_radii.get(element, 0.5)
                self.atom_radii.append(radius)

                coordReader = self.add_processor('envision.CoordinateReader', '{0} {1}'.format(i,name), xpos-i*7, ypos)
                self.network.addConnection(hdf5_output, coordReader.getInport('inport'))
                self.network.addConnection(coordReader.getOutport('outport'), strucMesh.getInport('coordinates'))
                coordReader.path.value = base_group + '/Atoms/' + key
                strucMesh_radius_property = strucMesh.getPropertyByIdentifier('radius{0}'.format(i))
                # The atoms in a crystal don't actually look like spheres, as the valence electrons are shared across the crystal.
                # The different radii of the elements in data.py are just to differentiate between different elements.
                strucMesh_radius_property.maxValue = 10
                strucMesh_radius_property.value = radius
                strucMesh_color_property = strucMesh.getPropertyByIdentifier('color{0}'.format(i))
                strucMesh_color_property.value = inviwopy.glm.vec4(color[0],color[1],color[2],color[3])

                strucMesh_atom_property = strucMesh.getPropertyByIdentifier('atoms{0}'.format(i))
                strucMesh_atom_property.value = 0
                strucMesh_atom_property.minValue = 0
                strucMesh_atom_property.maxValue = 0

                self.nAtomTypes += 1

        self.decoration_outport = meshRenderer.getOutport('image')
        # self.decoration_inport =
        self.decoration_inport = meshRenderer.getInport('imageInport')
        self.camera_prop = meshRenderer.camera
<commit_msg>atom works for ui<commit_after>import inviwopy
import numpy as np
import h5py
from envisionpy.utils.exceptions import *
from envisionpy.utils.atomData import atomic_radii, element_names, element_colors
from .baseNetworks.Decoration import Decoration

class AtomPositions(Decoration):
    '''
    Manages a subnetwork for atom position rendering.
    '''
    def __init__(self, inviwoApp, hdf5_path, hdf5_output, xpos=0, ypos=0):
        Decoration.__init__(self, inviwoApp)
        self.atom_radii = []
        self.atom_names = []
        self.nAtomTypes = 0
        self.setup_network(hdf5_path, hdf5_output, xpos, ypos)

        self.set_atom_radius(0.5)
        self.hide()

    @staticmethod
    def valid_hdf5(hdf5_file):
        return hdf5_file.get("UnitCell") != None

    def valid_visualisations(self):
        return ['charge', 'elf', 'parchg']

    def connect_decoration(self, other, vis_type):
        # Add a decoration by connecting data ports and linking properties.
        if vis_type not in self.valid_visualisations():
            raise EnvisionError('Invalid visualisation type ['+vis_type+'].')

        # Link needed properties between networks.
        if vis_type == 'charge' or vis_type == 'elf' or vis_type == 'parchg':
            self.network.addLink(self.camera_prop, other.camera_prop)
            self.network.addLink(other.camera_prop, self.camera_prop)

        other.connect_decoration_ports(self.decoration_outport)

    def disconnect_decoration(self, other, vis_type):
        if vis_type == 'charge' or vis_type == 'elf' or vis_type == 'parchg':
            self.network.removeLink(self.camera_prop, other.camera_prop)
        other.disconnect_decorations_port(self.decoration_outport)

    def show(self):
        self.get_processor('UnitcellCanvas').widget.show()

    def hide(self):
        self.get_processor('UnitcellCanvas').widget.hide()

    def get_ui_data(self):
        return [
            self.atom_names,
            self.atom_radii
        ]
# ------------------------------------------
# ------- Property control functions -------

    def set_atom_radius(self, radius, index=None):
        structureMesh = self.get_processor('UnitcellMesh')
        if structureMesh.fullMesh.value:
            if index != None:
                self.atom_radii[index] = radius
                structureMesh.getPropertyByIdentifier("radius" + str(index)).value = radius
            else:
                for i in range(self.nAtomTypes):
                    self.atom_radii[i] = radius
                    self.set_atom_radius(radius, i)
        else:
            sphereRenderer = self.get_processor('UnitcellRenderer')
            sphereRenderer.sphereProperties.defaultRadius.value = radius
            for i in range(self.nAtomTypes):
                self.atom_radii[i] = radius

    def hide_atoms(self):
        return self.set_atom_radius(0)

    def toggle_full_mesh(self, enable):
        structMesh = self.get_processor('UnitcellMesh')
        structMesh.fullMesh.value = enable

    def set_canvas_position(self, x, y):
    # Updates the position of the canvas
    # Upper left corner will be at coordinate (x, y)
        unitcellCanvas = self.get_processor('UnitcellCanvas')
        unitcellCanvas.position.value = inviwopy.glm.ivec2(x, y)

# ------------------------------------------
# ------- Network building functions -------

    def setup_network(self, hdf5_path, hdf5_output, xpos, ypos):
        strucMesh = self.add_processor('envision.StructureMesh', 'UnitcellMesh', xpos, ypos+3)
        meshRenderer = self.add_processor('org.inviwo.SphereRenderer', 'UnitcellRenderer', xpos, ypos+6)
        background = self.add_processor('org.inviwo.Background', 'AtomBackground', xpos, ypos+9)
        canvas = self.add_processor('org.inviwo.CanvasGL', 'UnitcellCanvas', xpos, ypos+12)

        canvas.inputSize.dimensions.value = inviwopy.glm.size2_t(500, 500)

        self.network.addConnection(strucMesh.getOutport('mesh'), meshRenderer.getInport('geometry'))
        self.network.addConnection(meshRenderer.getOutport('image'), background.getInport('inport'))
        self.network.addConnection(background.getOutport('outport'), canvas.getInport('inport'))

        with h5py.File(hdf5_path, "r") as h5:
            # Set basis matrix and scaling
            basis_matrix = np.array(h5["/basis"], dtype='d')
            strucMesh.basis.minValue = inviwopy.glm.mat3(
                -1000, -1000, -1000,
                -1000, -1000, -1000,
                -1000, -1000, -1000)
            strucMesh.basis.maxValue = inviwopy.glm.mat3(
                1000,1000,1000,
                1000,1000,1000,
                1000,1000,1000)
            strucMesh.basis.value = inviwopy.glm.mat3(
                basis_matrix[0,0], basis_matrix[0,1], basis_matrix[0,2],
                basis_matrix[1,0], basis_matrix[1,1], basis_matrix[1,2],
                basis_matrix[2,0], basis_matrix[2,1], basis_matrix[2,2])
            strucMesh.scalingFactor.maxValue = h5['/scaling_factor'][()]
            strucMesh.scalingFactor.value = h5['/scaling_factor'][()]

            # Setup animation?
            strucMesh.timestep.value = 0
            strucMesh.timestep.minValue = 0
            base_group = "/UnitCell"

            strucMesh.timestep.maxValue = 0
            for i,key in enumerate(list(h5[base_group + "/Atoms"].keys())):

                element = h5[base_group + "/Atoms/"+key].attrs['element']

                name = element_names.get(element, 'Unknown')
                self.atom_names.append(name)
                color = element_colors.get(element, (0.5, 0.5, 0.5, 1.0))
                radius = atomic_radii.get(element, 0.5)
                self.atom_radii.append(radius)

                coordReader = self.add_processor('envision.CoordinateReader', '{0} {1}'.format(i,name), xpos-i*7, ypos)
                self.network.addConnection(hdf5_output, coordReader.getInport('inport'))
                self.network.addConnection(coordReader.getOutport('outport'), strucMesh.getInport('coordinates'))
                coordReader.path.value = base_group + '/Atoms/' + key
                strucMesh_radius_property = strucMesh.getPropertyByIdentifier('radius{0}'.format(i))
                # The atoms in a crystal don't actually look like spheres, as the valence electrons are shared across the crystal.
                # The different radii of the elements in data.py are just to differentiate between different elements.
                strucMesh_radius_property.maxValue = 10
                strucMesh_radius_property.value = radius
                strucMesh_color_property = strucMesh.getPropertyByIdentifier('color{0}'.format(i))
                strucMesh_color_property.value = inviwopy.glm.vec4(color[0],color[1],color[2],color[3])

                strucMesh_atom_property = strucMesh.getPropertyByIdentifier('atoms{0}'.format(i))
                strucMesh_atom_property.value = 0
                strucMesh_atom_property.minValue = 0
                strucMesh_atom_property.maxValue = 0

                self.nAtomTypes += 1

        self.decoration_outport = meshRenderer.getOutport('image')
        # self.decoration_inport =
        self.decoration_inport = meshRenderer.getInport('imageInport')
        self.camera_prop = meshRenderer.camera
<|endoftext|>


<commit_before>"""
This example is largely adapted from https://github.com/pytorch/examples/blob/master/imagenet/main.py

Before you can run this example, you will need to download the ImageNet dataset manually from the
`official website <http://image-net.org/download>`_ and place it into a folder `path/to/imagenet`.

Train on ImageNet with default parameters:

.. code-block: bash

    python imagenet.py --data-path /path/to/imagenet

or show all options you can change:

.. code-block: bash

    python imagenet.py --help

"""
import os
from argparse import ArgumentParser, Namespace
from collections import OrderedDict

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

import pytorch_lightning as pl
from pytorch_lightning.core import LightningModule

class ImageNetLightningModel(LightningModule):

    # pull out resnet names from torchvision models
    MODEL_NAMES = sorted(
        name for name in models.__dict__
        if name.islower() and not name.startswith("__") and callable(models.__dict__[name])
    )

    def __init__(
        self,
        arch: str,
        pretrained: bool,
        lr: float,
        momentum: float,
        weight_decay: int,
        data_path: str,
        batch_size: int,
        workers: int,
        **kwargs,
    ):
        super().__init__()
        self.save_hyperparameters()
        self.arch = arch
        self.pretrained = pretrained
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.data_path = data_path
        self.batch_size = batch_size
        self.workers = workers
        self.model = models.__dict__[self.arch](pretrained=self.pretrained)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images, target = batch
        output = self(images)
        loss_val = F.cross_entropy(output, target)
        acc1, acc5 = self.__accuracy(output, target, topk=(1, 5))

        tqdm_dict = {'train_loss': loss_val}
        output = OrderedDict({
            'loss': loss_val,
            'acc1': acc1,
            'acc5': acc5,
            'progress_bar': tqdm_dict,
            'log': tqdm_dict
        })
        return output

    def validation_step(self, batch, batch_idx):
        images, target = batch
        output = self(images)
        loss_val = F.cross_entropy(output, target)
        acc1, acc5 = self.__accuracy(output, target, topk=(1, 5))

        output = OrderedDict({
            'val_loss': loss_val,
            'val_acc1': acc1,
            'val_acc5': acc5,
        })
        return output

    def validation_epoch_end(self, outputs):
        tqdm_dict = {}
        for metric_name in ["val_loss", "val_acc1", "val_acc5"]:
            tqdm_dict[metric_name] = torch.stack([output[metric_name] for output in outputs]).mean()

        result = {'progress_bar': tqdm_dict, 'log': tqdm_dict, 'val_loss': tqdm_dict["val_loss"]}
        return result

    @staticmethod
    def __accuracy(output, target, topk=(1,)):
        """Computes the accuracy over the k top predictions for the specified values of k"""
        with torch.no_grad():
            maxk = max(topk)
            batch_size = target.size(0)

            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))

            res = []
            for k in topk:
                correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100.0 / batch_size))
            return res

    def configure_optimizers(self):
        optimizer = optim.SGD(
            self.parameters(),
            lr=self.lr,
            momentum=self.momentum,
            weight_decay=self.weight_decay
        )
        scheduler = lr_scheduler.LambdaLR(
            optimizer,
            lambda epoch: 0.1 ** (epoch // 30)
        )
        return [optimizer], [scheduler]

    def train_dataloader(self):
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        )

        train_dir = os.path.join(self.data_path, 'train')
        train_dataset = datasets.ImageFolder(
            train_dir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))

        train_loader = torch.utils.data.DataLoader(
            dataset=train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.workers,
        )
        return train_loader

    def val_dataloader(self):
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        )
        val_dir = os.path.join(self.data_path, 'val')
        val_loader = torch.utils.data.DataLoader(
            datasets.ImageFolder(val_dir, transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ])),
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.workers,
        )
        return val_loader

    def test_dataloader(self):
        return self.val_dataloader()

    def test_step(self, *args, **kwargs):
        return self.validation_step(*args, **kwargs)

    def test_epoch_end(self, *args, **kwargs):
        outputs = self.validation_epoch_end(*args, **kwargs)

        def substitute_val_keys(out):
            return {k.replace('val', 'test'): v for k, v in out.items()}

        outputs = {
            'test_loss': outputs['val_loss'],
            'progress_bar': substitute_val_keys(outputs['progress_bar']),
            'log': substitute_val_keys(outputs['log']),
        }
        return outputs

    @staticmethod
    def add_model_specific_args(parent_parser):  # pragma: no-cover
        parser = ArgumentParser(parents=[parent_parser])
        parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
                            choices=ImageNetLightningModel.MODEL_NAMES,
                            help=('model architecture: ' + ' | '.join(ImageNetLightningModel.MODEL_NAMES)
                                  + ' (default: resnet18)'))
        parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                            help='number of data loading workers (default: 4)')
        parser.add_argument('-b', '--batch-size', default=256, type=int,
                            metavar='N',
                            help='mini-batch size (default: 256), this is the total '
                                 'batch size of all GPUs on the current node when '
                                 'using Data Parallel or Distributed Data Parallel')
        parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                            metavar='LR', help='initial learning rate', dest='lr')
        parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                            help='momentum')
        parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                            metavar='W', help='weight decay (default: 1e-4)',
                            dest='weight_decay')
        parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                            help='use pre-trained model')
        return parser

def main(args: Namespace) -> None:
    if args.seed is not None:
        pl.seed_everything(args.seed)

    if args.distributed_backend == 'ddp':
        # When using a single GPU per process and per
        # DistributedDataParallel, we need to divide the batch size
        # ourselves based on the total number of GPUs we have
        args.batch_size = int(args.batch_size / max(1, args.gpus))
        args.workers = int(args.workers / max(1, args.gpus))

    model = ImageNetLightningModel(**vars(args))
    trainer = pl.Trainer.from_argparse_args(args)

    if args.evaluate:
        trainer.test(model)
    else:
        trainer.fit(model)

def run_cli():
    parent_parser = ArgumentParser(add_help=False)
    parent_parser = pl.Trainer.add_argparse_args(parent_parser)
    parent_parser.add_argument('--data-path', metavar='DIR', type=str,
                               help='path to dataset')
    parent_parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                               help='evaluate model on validation set')
    parent_parser.add_argument('--seed', type=int, default=42,
                               help='seed for initializing training.')
    parser = ImageNetLightningModel.add_model_specific_args(parent_parser)
    parser.set_defaults(
        profiler=True,
        deterministic=True,
        max_epochs=90,
    )
    args = parser.parse_args()
    main(args)

if __name__ == '__main__':
    run_cli()
<commit_msg>Update to latest logging format and modify the accuracy method. (#4816)<commit_after>"""
This example is largely adapted from https://github.com/pytorch/examples/blob/master/imagenet/main.py

Before you can run this example, you will need to download the ImageNet dataset manually from the
`official website <http://image-net.org/download>`_ and place it into a folder `path/to/imagenet`.

Train on ImageNet with default parameters:

.. code-block: bash

    python imagenet.py --data-path /path/to/imagenet

or show all options you can change:

.. code-block: bash

    python imagenet.py --help

"""
import os
from argparse import ArgumentParser, Namespace

import torch
import torch.nn.functional as F
import torch.nn.parallel
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
import torch.utils.data
import torch.utils.data.distributed
import torchvision.datasets as datasets
import torchvision.models as models
import torchvision.transforms as transforms

import pytorch_lightning as pl
from pytorch_lightning.core import LightningModule

class ImageNetLightningModel(LightningModule):
    # pull out resnet names from torchvision models
    MODEL_NAMES = sorted(
        name for name in models.__dict__
        if name.islower() and not name.startswith("__") and callable(models.__dict__[name])
    )

    def __init__(
            self,
            arch: str,
            pretrained: bool,
            lr: float,
            momentum: float,
            weight_decay: int,
            data_path: str,
            batch_size: int,
            workers: int,
            **kwargs,
    ):
        super().__init__()
        self.save_hyperparameters()
        self.arch = arch
        self.pretrained = pretrained
        self.lr = lr
        self.momentum = momentum
        self.weight_decay = weight_decay
        self.data_path = data_path
        self.batch_size = batch_size
        self.workers = workers
        self.model = models.__dict__[self.arch](pretrained=self.pretrained)

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images, target = batch
        output = self(images)
        loss_train= F.cross_entropy(output, target)
        acc1, acc5 = self.__accuracy(output, target, topk=(1, 5))
        self.log('train_loss', loss_train, on_step=True, on_epoch=True, logger=True)
        self.log('train_acc1', acc1, on_step=True, prog_bar=True, on_epoch=True, logger=True)
        self.log('train_acc5', acc5, on_step=True, on_epoch=True, logger=True)
        return loss_val

    def validation_step(self, batch, batch_idx):
        images, target = batch
        output = self(images)
        loss_val = F.cross_entropy(output, target)
        acc1, acc5 = self.__accuracy(output, target, topk=(1, 5))
        self.log('val_loss', loss_val, on_step=True, on_epoch=True)
        self.log('val_acc1', acc1, on_step=True, prog_bar=True, on_epoch=True)
        self.log('val_acc5', acc5, on_step=True, on_epoch=True)

    @staticmethod
    def __accuracy(output, target, topk=(1,)):
        """Computes the accuracy over the k top predictions for the specified values of k"""
        with torch.no_grad():
            maxk = max(topk)
            batch_size = target.size(0)

            _, pred = output.topk(maxk, 1, True, True)
            pred = pred.t()
            correct = pred.eq(target.view(1, -1).expand_as(pred))

            res = []
            for k in topk:
                correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
                res.append(correct_k.mul_(100.0 / batch_size))
            return res

    def configure_optimizers(self):
        optimizer = optim.SGD(
            self.parameters(),
            lr=self.lr,
            momentum=self.momentum,
            weight_decay=self.weight_decay
        )
        scheduler = lr_scheduler.LambdaLR(
            optimizer,
            lambda epoch: 0.1 ** (epoch // 30)
        )
        return [optimizer], [scheduler]

    def train_dataloader(self):
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        )

        train_dir = os.path.join(self.data_path, 'train')
        train_dataset = datasets.ImageFolder(
            train_dir,
            transforms.Compose([
                transforms.RandomResizedCrop(224),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                normalize,
            ]))

        train_loader = torch.utils.data.DataLoader(
            dataset=train_dataset,
            batch_size=self.batch_size,
            shuffle=True,
            num_workers=self.workers,
        )
        return train_loader

    def val_dataloader(self):
        normalize = transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225],
        )
        val_dir = os.path.join(self.data_path, 'val')
        val_loader = torch.utils.data.DataLoader(
            datasets.ImageFolder(val_dir, transforms.Compose([
                transforms.Resize(256),
                transforms.CenterCrop(224),
                transforms.ToTensor(),
                normalize,
            ])),
            batch_size=self.batch_size,
            shuffle=False,
            num_workers=self.workers,
        )
        return val_loader

    def test_dataloader(self):
        return self.val_dataloader()

    def test_step(self, *args, **kwargs):
        return self.validation_step(*args, **kwargs)

    def test_epoch_end(self, *args, **kwargs):
        outputs = self.validation_epoch_end(*args, **kwargs)

        def substitute_val_keys(out):
            return {k.replace('val', 'test'): v for k, v in out.items()}

        outputs = {
            'test_loss': outputs['val_loss'],
            'progress_bar': substitute_val_keys(outputs['progress_bar']),
            'log': substitute_val_keys(outputs['log']),
        }
        return outputs

    @staticmethod
    def add_model_specific_args(parent_parser):  # pragma: no-cover
        parser = ArgumentParser(parents=[parent_parser])
        parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet18',
                            choices=ImageNetLightningModel.MODEL_NAMES,
                            help=('model architecture: ' + ' | '.join(ImageNetLightningModel.MODEL_NAMES)
                                  + ' (default: resnet18)'))
        parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                            help='number of data loading workers (default: 4)')
        parser.add_argument('-b', '--batch-size', default=256, type=int,
                            metavar='N',
                            help='mini-batch size (default: 256), this is the total '
                                 'batch size of all GPUs on the current node when '
                                 'using Data Parallel or Distributed Data Parallel')
        parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
                            metavar='LR', help='initial learning rate', dest='lr')
        parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                            help='momentum')
        parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
                            metavar='W', help='weight decay (default: 1e-4)',
                            dest='weight_decay')
        parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                            help='use pre-trained model')
        return parser

def main(args: Namespace) -> None:
    if args.seed is not None:
        pl.seed_everything(args.seed)

    if args.distributed_backend == 'ddp':
        # When using a single GPU per process and per
        # DistributedDataParallel, we need to divide the batch size
        # ourselves based on the total number of GPUs we have
        args.batch_size = int(args.batch_size / max(1, args.gpus))
        args.workers = int(args.workers / max(1, args.gpus))

    model = ImageNetLightningModel(**vars(args))
    trainer = pl.Trainer.from_argparse_args(args)

    if args.evaluate:
        trainer.test(model)
    else:
        trainer.fit(model)

def run_cli():
    parent_parser = ArgumentParser(add_help=False)
    parent_parser = pl.Trainer.add_argparse_args(parent_parser)
    parent_parser.add_argument('--data-path', metavar='DIR', type=str,
                               help='path to dataset')
    parent_parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                               help='evaluate model on validation set')
    parent_parser.add_argument('--seed', type=int, default=42,
                               help='seed for initializing training.')
    parser = ImageNetLightningModel.add_model_specific_args(parent_parser)
    parser.set_defaults(
        profiler=True,
        deterministic=True,
        max_epochs=90,
    )
    args = parser.parse_args()
    main(args)

if __name__ == '__main__':
    run_cli()
<|endoftext|>


<commit_before># Copyright 2022 The Scenic Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Data generators for the BAIR Robot dataset."""

import functools
from typing import Optional

from absl import logging
from dmvr import processors
from flax import jax_utils
import jax.numpy as jnp
from scenic.dataset_lib import dataset_utils
from scenic.dataset_lib import datasets
import tensorflow as tf

def preprocess_train_example(example, dtype=tf.float32, zero_centering=True):
  """Preprocesses the given video.

  Args:
    example: dict; Example that has an 'image_main'.
    dtype: Tensorflow data type; Data type of the image.
    zero_centering: If True, frames are normalized to values in [-1, 1].
      If False, values in [0, 1].

  Returns:
    dict; Example that has an 'inputs'.
  """
  frames = example['image_main']
  frames = processors.normalize_image(frames, zero_centering, dtype)
  return {'inputs': frames}

def augment_train_example(example, num_frames=30, stride=1):
  """Augment the given video for training.

  Args:
    example: dict; Example that has an 'inputs'.
    num_frames: Number of frames per subclip.
    stride: Temporal stride to sample frames.

  Returns:
    dict; Example that has an 'inputs'.
  """
  frames = example['inputs']
  frames = processors.sample_sequence(frames, num_frames, True, stride)
  frames = processors.random_flip_left_right(frames)
  return {'inputs': frames}

def preprocess_eval_example(example,
                            dtype=tf.float32,
                            num_frames=30,
                            stride=1,
                            num_clips=1,
                            zero_centering=True):
  """Preprocesses the given video for evaluation.

  Args:
    example: dict; Example that has an 'inputs'.
    dtype: Tensorflow data type; Data type of the image.
    num_frames: Number of frames per subclip.
    stride: Temporal stride to sample frames.
    num_clips: Linearly spaced clips to sample from each example.
    zero_centering: If True, frames are normalized to values in [-1, 1].
      If False, values in [0, 1].

  Returns:
    dict; Example that has an 'inputs'.
  """
  frames = example['image_main']
  frames = processors.normalize_image(frames, zero_centering, dtype)
  clips = processors.sample_linspace_sequence(frames, num_clips, num_frames,
                                              stride)
  return {'inputs': clips}

def postprocess_eval_batch(batch, num_frames=30):
  """Postprocesses the given batch for evaluation.

  Reshapes the batch from [bs, num_clips * num_frames, ...] into
    [bs * num_clips, num_frames, ...].

  Args:
    batch: dict; Batch that has an 'inputs'.
    num_frames: Number of frames per subclip.
  Returns:
    dict; Example that has an 'inputs'.
  """
  batch_clips = batch['inputs']
  batch_clips = tf.reshape(batch_clips,
                           (-1, num_frames, *batch_clips.shape[2:]))
  return {'inputs': batch_clips}

@datasets.add_dataset('bair')
def get_dataset(*,
                batch_size,
                eval_batch_size,
                num_shards,
                dtype_str='float32',
                shuffle_seed=0,
                rng=None,
                dataset_configs=None,
                dataset_service_address: Optional[str] = None):
  """Returns generators for the BAIR train, validation, and test set.

  Args:
    batch_size: int; Determines the train batch size.
    eval_batch_size: int; Determines the evaluation batch size.
    num_shards: int;  Number of shards --> batch shape: [num_shards, bs, ...].
    dtype_str: Data type of the image (e.g. 'float32').
    shuffle_seed: int; Seed for shuffling the training data.
    rng: JAX rng key, which can be used for augmentation, shuffling, etc.
    dataset_configs: dict; Dataset specific configurations.
    dataset_service_address: If set, will distribute the training dataset using
      the given tf.data service at the given address.

  Returns:
    A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
      a test_iter, and a dict of meta_data.
  """
  del rng
  dtype = getattr(tf, dtype_str)
  dataset_configs = dataset_configs or {}
  num_frames = dataset_configs.get('num_frames', 30)
  stride = dataset_configs.get('stride', 1)
  zero_centering = dataset_configs.get('zero_centering', True)
  num_eval_clips = dataset_configs.get('num_eval_clips', 1)
  shuffle_buffer_size = dataset_configs.get('shuffle_buffer_size', None)
  preprocess_train = functools.partial(
      preprocess_train_example, dtype=dtype, zero_centering=zero_centering)
  augment_train = functools.partial(
      augment_train_example, num_frames=num_frames, stride=stride)
  preprocess_eval = functools.partial(
      preprocess_eval_example,
      dtype=dtype,
      num_frames=num_frames,
      stride=stride,
      num_clips=num_eval_clips,
      zero_centering=zero_centering)
  if num_eval_clips > 1:
    postprocess_eval = functools.partial(
        postprocess_eval_batch, num_frames=num_frames)
  else:
    postprocess_eval = None

  logging.info('Loading train split of the BAIR dataset.')
  train_ds, _ = dataset_utils.load_split_from_tfds(
      'bair_robot_pushing_small',
      batch_size,
      split='train',
      preprocess_example=preprocess_train,
      augment_train_example=augment_train,
      shuffle_buffer_size=shuffle_buffer_size,
      shuffle_seed=shuffle_seed)

  if dataset_service_address:
    if shuffle_seed is not None:
      raise ValueError('Using dataset service with a random seed causes each '
                       'worker to produce exactly the same data. Add '
                       'config.shuffle_seed = None to your config if you '
                       'want to run with dataset service.')
    logging.info('Using the tf.data service at %s', dataset_service_address)
    train_ds = dataset_utils.distribute(train_ds, dataset_service_address)

  logging.info('Loading test split of the BAIR dataset.')
  eval_ds, _ = dataset_utils.load_split_from_tfds(
      'bair_robot_pushing_small',
      eval_batch_size,
      split='test',
      preprocess_example=preprocess_eval,
      postprocess_batch=postprocess_eval)

  maybe_pad_batches_train = functools.partial(
      dataset_utils.maybe_pad_batch, train=True, batch_size=batch_size)
  maybe_pad_batches_eval = functools.partial(
      dataset_utils.maybe_pad_batch,
      train=False,
      batch_size=eval_batch_size * num_eval_clips)
  shard_batches = functools.partial(dataset_utils.shard, n_devices=num_shards)

  train_iter = iter(train_ds)
  train_iter = map(dataset_utils.tf_to_numpy, train_iter)
  train_iter = map(maybe_pad_batches_train, train_iter)
  train_iter = map(shard_batches, train_iter)
  if dataset_configs.get('prefetch_to_device'):
    # Async bind batch to device which speeds up training.
    train_iter = jax_utils.prefetch_to_device(
        train_iter, dataset_configs.get('prefetch_to_device'))

  eval_iter = iter(eval_ds)
  eval_iter = map(dataset_utils.tf_to_numpy, eval_iter)
  eval_iter = map(maybe_pad_batches_eval, eval_iter)
  eval_iter = map(shard_batches, eval_iter)

  input_shape = (-1, num_frames, 64, 64, 3)
  num_train_examples = dataset_utils.get_num_examples(
      'bair_robot_pushing_small', 'train')
  num_eval_examples = dataset_utils.get_num_examples('bair_robot_pushing_small',
                                                     'test') * num_eval_clips
  meta_data = {
      'num_classes': None,
      'input_shape': input_shape,
      'num_train_examples': num_train_examples,
      'num_eval_examples': num_eval_examples,
      'input_dtype': getattr(jnp, dtype_str),
      'target_is_onehot': False,
  }
  return dataset_utils.Dataset(train_iter, eval_iter, None, meta_data)
<commit_msg>Add camera selection for bair dataset from tfds.<commit_after># Copyright 2022 The Scenic Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Data generators for the BAIR Robot dataset."""

import functools
from typing import Optional

from absl import logging
from dmvr import processors
from flax import jax_utils
import jax.numpy as jnp
from scenic.dataset_lib import dataset_utils
from scenic.dataset_lib import datasets
import tensorflow as tf

def preprocess_train_example(example,
                             camera_name='image_main',
                             dtype=tf.float32,
                             zero_centering=True):
  """Preprocesses the given video.

  Args:
    example: dict; Example that has an 'image_main'.
    camera_name: Name of the image sequence to use.
    dtype: Tensorflow data type; Data type of the image.
    zero_centering: If True, frames are normalized to values in [-1, 1].
      If False, values in [0, 1].

  Returns:
    dict; Example that has an 'inputs'.
  """
  frames = example[camera_name]
  frames = processors.normalize_image(frames, zero_centering, dtype)
  return {'inputs': frames}

def augment_train_example(example, num_frames=30, stride=1):
  """Augment the given video for training.

  Args:
    example: dict; Example that has an 'inputs'.
    num_frames: Number of frames per subclip.
    stride: Temporal stride to sample frames.

  Returns:
    dict; Example that has an 'inputs'.
  """
  frames = example['inputs']
  frames = processors.sample_sequence(frames, num_frames, True, stride)
  frames = processors.random_flip_left_right(frames)
  return {'inputs': frames}

def preprocess_eval_example(example,
                            camera_name='image_main',
                            dtype=tf.float32,
                            num_frames=30,
                            stride=1,
                            num_clips=1,
                            zero_centering=True):
  """Preprocesses the given video for evaluation.

  Args:
    example: dict; Example that has an 'inputs'.
    camera_name: Name of the image sequence to use.
    dtype: Tensorflow data type; Data type of the image.
    num_frames: Number of frames per subclip.
    stride: Temporal stride to sample frames.
    num_clips: Linearly spaced clips to sample from each example.
    zero_centering: If True, frames are normalized to values in [-1, 1].
      If False, values in [0, 1].

  Returns:
    dict; Example that has an 'inputs'.
  """
  frames = example[camera_name]
  frames = processors.normalize_image(frames, zero_centering, dtype)
  clips = processors.sample_linspace_sequence(frames, num_clips, num_frames,
                                              stride)
  return {'inputs': clips}

def postprocess_eval_batch(batch, num_frames=30):
  """Postprocesses the given batch for evaluation.

  Reshapes the batch from [bs, num_clips * num_frames, ...] into
    [bs * num_clips, num_frames, ...].

  Args:
    batch: dict; Batch that has an 'inputs'.
    num_frames: Number of frames per subclip.
  Returns:
    dict; Example that has an 'inputs'.
  """
  batch_clips = batch['inputs']
  batch_clips = tf.reshape(batch_clips,
                           (-1, num_frames, *batch_clips.shape[2:]))
  return {'inputs': batch_clips}

@datasets.add_dataset('bair')
def get_dataset(*,
                batch_size,
                eval_batch_size,
                num_shards,
                dtype_str='float32',
                shuffle_seed=0,
                rng=None,
                dataset_configs=None,
                dataset_service_address: Optional[str] = None):
  """Returns generators for the BAIR train, validation, and test set.

  Args:
    batch_size: int; Determines the train batch size.
    eval_batch_size: int; Determines the evaluation batch size.
    num_shards: int;  Number of shards --> batch shape: [num_shards, bs, ...].
    dtype_str: Data type of the image (e.g. 'float32').
    shuffle_seed: int; Seed for shuffling the training data.
    rng: JAX rng key, which can be used for augmentation, shuffling, etc.
    dataset_configs: dict; Dataset specific configurations.
    dataset_service_address: If set, will distribute the training dataset using
      the given tf.data service at the given address.

  Returns:
    A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
      a test_iter, and a dict of meta_data.
  """
  del rng
  dtype = getattr(tf, dtype_str)
  dataset_configs = dataset_configs or {}
  camera_name = dataset_configs.get('camera_name', 'image_main')
  num_frames = dataset_configs.get('num_frames', 30)
  stride = dataset_configs.get('stride', 1)
  zero_centering = dataset_configs.get('zero_centering', True)
  num_eval_clips = dataset_configs.get('num_eval_clips', 1)
  shuffle_buffer_size = dataset_configs.get('shuffle_buffer_size', None)
  preprocess_train = functools.partial(
      preprocess_train_example,
      camera_name=camera_name,
      dtype=dtype,
      zero_centering=zero_centering)
  augment_train = functools.partial(
      augment_train_example, num_frames=num_frames, stride=stride)
  preprocess_eval = functools.partial(
      preprocess_eval_example,
      camera_name=camera_name,
      dtype=dtype,
      num_frames=num_frames,
      stride=stride,
      num_clips=num_eval_clips,
      zero_centering=zero_centering)
  if num_eval_clips > 1:
    postprocess_eval = functools.partial(
        postprocess_eval_batch, num_frames=num_frames)
  else:
    postprocess_eval = None

  logging.info('Loading train split of the BAIR dataset.')
  train_ds, _ = dataset_utils.load_split_from_tfds(
      'bair_robot_pushing_small',
      batch_size,
      split='train',
      preprocess_example=preprocess_train,
      augment_train_example=augment_train,
      shuffle_buffer_size=shuffle_buffer_size,
      shuffle_seed=shuffle_seed)

  if dataset_service_address:
    if shuffle_seed is not None:
      raise ValueError('Using dataset service with a random seed causes each '
                       'worker to produce exactly the same data. Add '
                       'config.shuffle_seed = None to your config if you '
                       'want to run with dataset service.')
    logging.info('Using the tf.data service at %s', dataset_service_address)
    train_ds = dataset_utils.distribute(train_ds, dataset_service_address)

  logging.info('Loading test split of the BAIR dataset.')
  eval_ds, _ = dataset_utils.load_split_from_tfds(
      'bair_robot_pushing_small',
      eval_batch_size,
      split='test',
      preprocess_example=preprocess_eval,
      postprocess_batch=postprocess_eval)

  maybe_pad_batches_train = functools.partial(
      dataset_utils.maybe_pad_batch, train=True, batch_size=batch_size)
  maybe_pad_batches_eval = functools.partial(
      dataset_utils.maybe_pad_batch,
      train=False,
      batch_size=eval_batch_size * num_eval_clips)
  shard_batches = functools.partial(dataset_utils.shard, n_devices=num_shards)

  train_iter = iter(train_ds)
  train_iter = map(dataset_utils.tf_to_numpy, train_iter)
  train_iter = map(maybe_pad_batches_train, train_iter)
  train_iter = map(shard_batches, train_iter)
  if dataset_configs.get('prefetch_to_device'):
    # Async bind batch to device which speeds up training.
    train_iter = jax_utils.prefetch_to_device(
        train_iter, dataset_configs.get('prefetch_to_device'))

  eval_iter = iter(eval_ds)
  eval_iter = map(dataset_utils.tf_to_numpy, eval_iter)
  eval_iter = map(maybe_pad_batches_eval, eval_iter)
  eval_iter = map(shard_batches, eval_iter)

  input_shape = (-1, num_frames, 64, 64, 3)
  num_train_examples = dataset_utils.get_num_examples(
      'bair_robot_pushing_small', 'train')
  num_eval_examples = dataset_utils.get_num_examples('bair_robot_pushing_small',
                                                     'test') * num_eval_clips
  meta_data = {
      'num_classes': None,
      'input_shape': input_shape,
      'num_train_examples': num_train_examples,
      'num_eval_examples': num_eval_examples,
      'input_dtype': getattr(jnp, dtype_str),
      'target_is_onehot': False,
  }
  return dataset_utils.Dataset(train_iter, eval_iter, None, meta_data)
<|endoftext|>