def setup_integrator(method, use_adjoint, step_size, rtol=1e-3, atol=1e-5):
integrator_options = dict()
if method not in ['dopri5', 'adams']:
integrator_options = {'step_size': step_size}
integrator = generic_integrator.GenericIntegrator(
integrator_library = 'odeint',
integrator_name = method,
use_adjoint_integration=use_adjoint,
integrator_options=integrator_options,
rtol=rtol, atol=atol)
return integrator
def setup_integrator(method,
use_adjoint,
step_size,
rtol=1e-8,
atol=1e-12,
nr_of_checkpoints=None,
checkpointing_time_interval=None):
integrator_options = dict()
if method not in ['dopri5', 'adams']:
integrator_options = {'step_size': step_size}
integrator = generic_integrator.GenericIntegrator(
integrator_library='odeint',
integrator_name=method,
use_adjoint_integration=use_adjoint,
integrator_options=integrator_options,
rtol=rtol,
atol=atol,
nr_of_checkpoints=nr_of_checkpoints,
checkpointing_time_interval=checkpointing_time_interval)
return integrator
def __init__(self,
nr_of_image_channels,
layer_channels=[64, 128, 256, 512],
nr_of_blocks_per_layer=[1, 1, 1, 1],
downsampling_stride=2,
nonlinearity='relu',
particle_sizes=[[15, 15], [11, 11], [7, 7], [5, 5]],
nr_of_particles=10,
nr_of_classes=10,
state_initializer=None,
costate_initializer=None,
integrator=None,
shooting_model=None,
shooting_model_kwargs=None):
super(BasicResNet, self).__init__()
self.nr_of_image_channels = nr_of_image_channels
self.layer_channels = layer_channels
self.nr_of_blocks_per_layer = nr_of_blocks_per_layer
self.nr_of_classes = nr_of_classes
self.nonlinearity = nonlinearity
self.nl, _ = ad.get_nonlinearity(nonlinearity=nonlinearity)
if state_initializer is not None:
print('INFO: Using externally defined state initializer')
self._state_initializer = state_initializer
else:
print('WARNING: Using default state initializer')
self._state_initializer = parameter_initialization.ConvolutionEvolutionParameterInitializer(
only_random_initialization=True,
random_initialization_magnitude=0.05)
if costate_initializer is not None:
print('INFO: Using externally defined costate initializer')
self._costate_initializer = costate_initializer
else:
print('WARNING: Using default costate initializer')
self._costate_initializer = parameter_initialization.ConvolutionEvolutionParameterInitializer(
only_random_initialization=True,
random_initialization_magnitude=0.05)
# setup the integrator
if integrator is not None:
print('INFO: Using externally defined integrator')
self._integrator = integrator
else:
print('WARNING: Using default integrator')
integrator_options = dict()
integrator_options['step_size'] = 0.1
self._integrator = generic_integrator.GenericIntegrator(
integrator_library='odeint',
integrator_name='rk4',
integrator_options=integrator_options)
self.nr_of_particles = nr_of_particles
self.particle_sizes = particle_sizes
if len(layer_channels) != len(self.particle_sizes):
raise ValueError(
'Dimension mismatch, between laters and particle sizes. A particle size needs to be defined for each layer.'
)
# initial convolution layer
self.initial_conv = nn.Conv2d(self.nr_of_image_channels,
layer_channels[0],
kernel_size=3,
padding=1,
bias=True)
self.batch_norm = nn.BatchNorm2d(layer_channels[0])
downsampling_strides = [downsampling_stride
] * (len(layer_channels) - 1) + [None]
if shooting_model is not None:
self.shooting_model = shooting_model
else:
self.shooting_model = shooting_models.AutoShootingIntegrandModelSimpleConv2D
print('WARNING: Using default shooting model {}'.format(
self.shooting_model))
if shooting_model_kwargs is not None:
self._shooting_model_kwargs = shooting_model_kwargs
else:
self._shooting_model_kwargs = {}
self.shooting_layers = self._create_shooting_layers(
layer_channels=self.layer_channels,
nr_of_blocks_per_layer=self.nr_of_blocks_per_layer)
self.striding_blocks = self._create_striding_blocks(
strides=downsampling_strides)
self.last_linear = nn.Linear(layer_channels[-1], self.nr_of_classes)
self._forward_not_yet_executed = True
def __init__(self,
name,
shooting_integrand=None,
shooting_integrand_name=None,
integrator_library='odeint',
integrator_name='rk4',
use_adjoint_integration=False,
integrator_options=None,
integrator=None,
keep_initial_state_parameters_at_zero=False,
enlarge_pass_through_states_and_costates=True,
use_particle_free_rnn_mode=False,
use_particle_free_time_dependent_mode=False,
nr_of_particle_free_time_dependent_steps=10,
max_integration_time=None,
*args,
**kwargs):
"""
:param name: unique name of this block (needed to keep track of the parameters when there is pass through)
:param shooting_integrand:
:param shooting_integrand_name:
:param integrator_library:
:param integrator_name:
:param use_adjoint_integration:
:param integrator_options:
:param integrator: We can also instantiate an integrator externally and simply pass it here
:param keep_initial_state_parameters_at_zero: If set to true than all the newly created initial state parameters are kept at zero (and not optimized over); this includes state parameters created via state/costate enlargement.
:param enlarge_pass_through_states_and_costates: all the pass through states/costates are enlarged so they match the dimensions of the states/costates. This assures that parameters can be concatenated.
:param use_particle_free_rnn_mode: if set to true than the particles are not used to compute the parameterization, instead an RNN model is assumed and the layer parameters are optimized directly, particles will stay as initialized
:param use_particle_free_time_dependent_mode: mimicks a more classical resnet, by parameterizing via a finite number of direct transformations
:param nr_of_particle_free_time_dependent_steps: specifies the number of steps where the parameterization happens (piecewise constant in the middle)
:param max_integration_time: sets the maximum integration time. So default integrations are from 0 to max_integration_times. Also affects time-interval for particle_free_time_dependent_mode
:param args:
:param kwargs:
"""
super(ShootingBlockBase, self).__init__()
self._block_name = name
"""Name of the shooting block"""
self._forward_not_yet_executed = True
"""This is to keep track if it has been run forward. To assure parameters() is not called before one forward pass has been accomplished.
As for convenience the parameters are constructed dynamically."""
if shooting_integrand is None and shooting_integrand_name is None:
raise ValueError(
'Either shooting_integrand or shooting_integrand_name need to be specified'
)
if shooting_integrand is not None and shooting_integrand_name is not None:
raise ValueError(
'Either shooting_integrand or shooting_integrand_name need to be specified. You specified both. Pick one option.'
)
if shooting_integrand is not None:
self.shooting_integrand = shooting_integrand
else:
self.shooting_integrand = self._get_shooting_integrand_by_name(
shooting_integrand_name=shooting_integrand_name)
#self.add_module(name='shooting_integrand', module=self.shooting_integrand) # probably remove
if max_integration_time is not None:
self.integration_time = torch.tensor([0, max_integration_time
]).float()
else:
# set to default integration interval
self.integration_time = torch.tensor([0, 1]).float()
self.integration_time_vector = None
if integrator is not None:
print('INFO: using externally specified integrator')
self.integrator = integrator
else:
self.integrator_name = integrator_name
self.integrator_library = integrator_library
self.integrator_options = integrator_options
self.use_adjoint_integration = use_adjoint_integration
if self.integrator_options is None:
print(
'WARNING: integrator options are NOT set. Will use default options'
)
self.integrator = generic_integrator.GenericIntegrator(
integrator_library=self.integrator_library,
integrator_name=self.integrator_name,
use_adjoint_integration=self.use_adjoint_integration,
integrator_options=self.integrator_options)
self._state_parameter_dict = None
"""Dictionary holding the state variables"""
self.keep_state_parameters_at_zero = keep_initial_state_parameters_at_zero
"""
If set to true one only optimizes over the costate and the state parameters are kept at zero (i.e., are no parameters).
This is for example useful when mimicking ResNet style dimension increase.
"""
self.enlarge_pass_through_states_and_costates = enlarge_pass_through_states_and_costates
"""If set to true the pass through states and costates will be enlarged so they are compatible in size with the states and costates and can be concatenated"""
self.use_particle_free_rnn_mode = use_particle_free_rnn_mode
"""if set to true than the particles are not used to compute the parameterization, instead an RNN model is assumed and the layer parameters are optimized directly, particles will stay as initialized"""
self._particle_free_rnn_parameter_func = None
self.use_particle_free_time_dependent_mode = use_particle_free_time_dependent_mode
self.nr_of_particle_free_time_dependent_steps = nr_of_particle_free_time_dependent_steps
if self.use_particle_free_time_dependent_mode and self.use_particle_free_rnn_mode:
raise ValueError(
'use_particle_free_time_depdenent_mode and use_particle_free_rnn_mode cannot be set simultaneosusly'
)
self._costate_parameter_dict = None
"""Dictionary holding the costates (i.e., adjoints/duals)"""
# TODO: see if we really need these variables, seems like they are no longer needed
# self._pass_through_state_parameter_dict_of_dicts = None
# """State parameters that are passed in externally, but are not parameters to be optimized over"""
# self._pass_through_costate_parameter_dict_of_dicts = None
# """Costate parameters that are passed in externally, but are not parameters to be optimized over"""
self._pass_through_state_dict_of_dicts_enlargement_parameters = None
self._pass_through_costate_dict_of_dicts_enlargement_parameters = None
state_dict, costate_dict = self.create_initial_state_and_costate_parameters(
*args, **kwargs)
# particle-free RNN mode is a direct optimization over the NN parameters (a la Neural ODE)
# it is included here simply to allow for direct comparisions with the same code
# if self.use_particle_free_rnn_mode:
# # those will have the default initialization
# self._particle_free_rnn_parameters = self.shooting_integrand.create_default_parameter_objects_on_consistent_device()
# # let's register them for optimization
# self._particle_free_rnn_parameters = self.register_particle_free_rnn_parameters(rnn_parameters=self._particle_free_rnn_parameters)
# # and then associate them with the integrand, so the integrand knows that these are externally managed
# self.shooting_integrand.set_externally_managed_rnn_parameters(self._particle_free_rnn_parameters)
# # we just set them to default, these will not really be needed though (as we are not optimizing over these values)
# self._state_parameter_dict = state_dict
# self._costate_parameter_dict = costate_dict
if self.use_particle_free_rnn_mode or self.use_particle_free_time_dependent_mode:
if self.use_particle_free_rnn_mode:
time_dependent_steps = 1
else:
time_dependent_steps = self.nr_of_particle_free_time_dependent_steps
# create the function that contains the parameters and that we can query for a parameter at a particular time
self._particle_free_rnn_parameter_func = ParticleFreeParameterFunc(
time_dependent_steps=time_dependent_steps,
integration_time=self.integration_time,
parameter_creation_func=self.shooting_integrand.
create_default_parameter_objects_on_consistent_device)
# register its parameters
pars = self._particle_free_rnn_parameter_func.get_time_parameter_list(
)
for i, p in enumerate(pars):
self.register_particle_free_parameters(
parameters=p, name_prefix='pf_timestep{}_'.format(i))
# and then associate them with the integrand, so the integrand knows that these are externally managed
self.shooting_integrand.set_externally_managed_rnn_parameter_func(
self._particle_free_rnn_parameter_func)
# we just set them to default, these will not really be needed though (as we are not optimizing over these values)
self._state_parameter_dict = state_dict
self._costate_parameter_dict = costate_dict
else:
self._state_parameter_dict, self._costate_parameter_dict = self.register_state_and_costate_parameters(
state_dict=state_dict,
costate_dict=costate_dict,
keep_state_parameters_at_zero=
keep_initial_state_parameters_at_zero)
# now register the parameters that are part of the integrands (this allows for easy definition of parameters in the integrand constructors
self.register_integrand_parameters(integrand=self.shooting_integrand)
in_features_size = 2
#check_models = ['updown']
#check_models = ['DEBUG']
check_models = ['simple']
#check_models = ['universal']
#check_models = ['updown','DEBUG','simple',"universal"]
number_of_tests_passed = 0
number_of_tests_attempted = 0
tolerance = 5e-3
integrator = generic_integrator.GenericIntegrator(
integrator_library='odeint',
integrator_name='rk4',
use_adjoint_integration=False,
integrator_options=integrator_options)
for current_model in check_models:
if current_model == 'simple':
shooting_model = shooting_models.AutoShootingIntegrandModelSimple(
in_features=in_features_size,
nonlinearity=nonlinearity,
nr_of_particles=nr_of_particles,
particle_dimension=1,
particle_size=in_features_size,
parameter_weight=parameter_weight)
elif current_model == 'universal':
shooting_model = shooting_models.AutoShootingIntegrandModelUniversal(
weight_decay = 0.0001
print('Using ResNet: weight = {}'.format(weight_decay))
simple_resnet = sdnn.ResNet(nr_of_layers=args.nr_of_layers)
elif args.use_double_resnet:
#weight_decay = 0.025
weight_decay = 0.01
print('Using DoubleResNetUpDown: weight = {}'.format(weight_decay))
simple_resnet = sdnn.DoubleResNetUpDown(
nr_of_layers=args.nr_of_layers, inflation_factor=inflation_factor)
elif args.use_neural_ode:
print('Using neural ode')
func = sdnn.ODESimpleFunc()
optimizer = optim.RMSprop(func.parameters(), lr=1e-3)
integrator = gi.GenericIntegrator(integrator_library='odeint',
integrator_name=args.method,
use_adjoint_integration=False,
rtol=1e-8,
atol=1e-12)
else:
use_shooting = True
print('Using shooting')
if not use_shooting:
if args.use_neural_ode:
params = func.parameters()
else:
params = simple_resnet.parameters()
else:
sample_batch_in, sample_batch_out = get_sample_batch(