def __init__( self, in_features: int, out_features: int, bias: bool = True, rpu_config: Optional[RPUConfigAlias] = None, realistic_read_write: bool = False, weight_scaling_omega: Optional[float] = None, ): # Call super() after tile creation, including ``reset_parameters``. Linear.__init__(self, in_features, out_features, bias=bias) # Create tiles if rpu_config is None: rpu_config = SingleRPUConfig() AnalogModuleBase.__init__( self, in_features, out_features, bias, realistic_read_write, rpu_config.mapping ) if self.analog_bias: raise ModuleError("AnalogLinearMapped only supports digital bias.") # More than one tile may need to be created. If so, divide # weight matrix into equal pieces along input dimension with # as many tiles as needed max_input_size = rpu_config.mapping.max_input_size max_output_size = rpu_config.mapping.max_output_size self.in_sizes = self.get_split_sizes(in_features, max_input_size) self.out_sizes = self.get_split_sizes(out_features, max_output_size) self.analog_tile_array = [] for i, in_tile_size in enumerate(self.in_sizes): in_tiles = [] for j, out_tile_size in enumerate(self.out_sizes): tile = rpu_config.tile_class(out_tile_size, in_tile_size, rpu_config, bias=self.analog_bias) self.register_analog_tile(tile, name=f"{i}_{j}") in_tiles.append(tile) self.analog_tile_array.append(in_tiles) # Set weights from the reset_parameters self.set_weights(self.weight, self.bias, remap_weights=True, weight_scaling_omega=weight_scaling_omega) # Unregister weight/bias as a parameter but keep for sync self.unregister_parameter('weight') if self.analog_bias: self.unregister_parameter('bias')
def init_analog_layer(layer: AnalogModuleBase) -> None: """Init the weights and bias of an analog linear layer.""" weight_init_fn(layer.weight.data) if layer.use_bias: if bias_init_fn is None: weight_init_fn(layer.bias.data) else: bias_init_fn(layer.bias.data) layer.set_weights(layer.weight, layer.bias)
def __init__( self, in_channels: int, out_channels: int, kernel_size: Tuple[int, ...], stride: Tuple[int, ...], padding: Tuple[int, ...], dilation: Tuple[int, ...], transposed: bool, output_padding: Tuple[int, ...], groups: int, bias: bool, padding_mode: str, rpu_config: Optional[RPUConfigAlias] = None, realistic_read_write: bool = False, weight_scaling_omega: Optional[float] = None, ): # pylint: disable=too-many-arguments, too-many-locals if groups != 1: raise ValueError('Only one group is supported') if padding_mode != 'zeros': raise ValueError('Only "zeros" padding mode is supported') # Call super() after tile creation, including ``reset_parameters``. _ConvNd.__init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding, groups, bias, padding_mode) # Create the tile and set the analog. if rpu_config is None: rpu_config = SingleRPUConfig() AnalogModuleBase.__init__( self, self.get_tile_size(in_channels, groups, kernel_size), out_channels, bias, realistic_read_write, rpu_config.mapping) self.analog_tile = self._setup_tile(rpu_config) # Register analog tile self.register_analog_tile(self.analog_tile) # Set weights from the reset_parameters self.set_weights(self.weight, self.bias, remap_weights=True, weight_scaling_omega=weight_scaling_omega) # Set the index matrices. self.fold_indices = Tensor().detach() self.input_size = 0 self.tensor_view = (-1, ) # type: Tuple[int, ...] # Unregister weight/bias as a parameter but keep it for syncs self.unregister_parameter('weight') if self.analog_bias: self.unregister_parameter('bias')
def extra_repr(self) -> str: """Set the extra representation of the module. Returns: A string with the extra representation. """ output = AnalogModuleBase.extra_repr(self)[:-1] output += ', mapping={}'.format((len(self.in_sizes), len(self.out_sizes))) return output
def __init__( self, in_features: int, out_features: int, bias: bool = True, rpu_config: Optional[RPUConfigAlias] = None, realistic_read_write: bool = False, weight_scaling_omega: Optional[float] = None, ): # Call super() after tile creation, including ``reset_parameters``. Linear.__init__(self, in_features, out_features, bias=bias) # Create tile if rpu_config is None: rpu_config = SingleRPUConfig() AnalogModuleBase.__init__( self, in_features, out_features, bias, realistic_read_write, weight_scaling_omega, rpu_config.mapping ) self.analog_tile = self._setup_tile(rpu_config) # Register tile self.register_analog_tile(self.analog_tile) # Set weights from the reset_parameters call self.set_weights(self.weight, self.bias) # Unregister weight/bias as a parameter but keep it as a # field (needed for syncing still) self.unregister_parameter('weight') if self.analog_bias: self.unregister_parameter('bias')
def __init__( self, in_channels: int, out_channels: int, kernel_size: Tuple[int, ...], stride: Tuple[int, ...], padding: Tuple[int, ...], dilation: Tuple[int, ...], transposed: bool, output_padding: Tuple[int, ...], groups: int, bias: bool, padding_mode: str, rpu_config: Optional[RPUConfigAlias] = None, realistic_read_write: bool = False, weight_scaling_omega: Optional[float] = None, ): # pylint: disable=too-many-arguments, too-many-locals if groups != 1: raise ValueError('Only one group is supported') if padding_mode != 'zeros': raise ValueError('Only "zeros" padding mode is supported') # Call super() after tile creation, including ``reset_parameters``. _ConvNd.__init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding, groups, bias, padding_mode) # Create tiles if rpu_config is None: rpu_config = SingleRPUConfig() AnalogModuleBase.__init__( self, self.get_tile_size(in_channels, groups, kernel_size), out_channels, bias, realistic_read_write, weight_scaling_omega, rpu_config.mapping) if self.analog_bias: raise ModuleError("AnalogConvNdMapped only supports digital bias.") if not rpu_config: rpu_config = SingleRPUConfig() max_input_size = rpu_config.mapping.max_input_size max_output_size = rpu_config.mapping.max_output_size kernel_elem = self.in_features // self.in_channels self.in_sizes = self.get_split_sizes(self.in_features, max_input_size, kernel_elem) self.out_sizes = self.get_split_sizes(self.out_features, max_output_size) self.analog_tile_array = [] for i, in_tile_size in enumerate(self.in_sizes): in_tiles = [] for j, out_tile_size in enumerate(self.out_sizes): tile = rpu_config.tile_class(out_tile_size, in_tile_size * kernel_elem, rpu_config, bias=self.analog_bias) self.register_analog_tile(tile, name=f"{i}_{j}") in_tiles.append(tile) self.analog_tile_array.append(in_tiles) # Set weights from the reset_parameters (since now the # analog_tiles are registered) self.set_weights(self.weight, self.bias) # Set the index matrices. self.input_size = 0 self.fold_indices_lst = [] # type: List[Tensor] # Unregister weight/bias as a parameter but keep it as a # field (needed for syncing still) self.unregister_parameter('weight') if self.analog_bias: self.unregister_parameter('bias')