def create_analog_network(): """Returns a Vgg8 inspired analog model.""" channel_base = 48 channel = [channel_base, 2 * channel_base, 3 * channel_base] fc_size = 8 * channel_base model = AnalogSequential( nn.Conv2d(in_channels=3, out_channels=channel[0], kernel_size=3, stride=1, padding=1), nn.ReLU(), AnalogConv2d(in_channels=channel[0], out_channels=channel[0], kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(channel[0]), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=channel[1], out_channels=channel[1], kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(channel[1]), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=channel[1], out_channels=channel[2], kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=channel[2], out_channels=channel[2], kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(channel[2]), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), nn.Flatten(), AnalogLinear(in_features=16 * channel[2], out_features=fc_size, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), nn.Linear(in_features=fc_size, out_features=N_CLASSES), nn.LogSoftmax(dim=1)) return model
def __init__(self): super().__init__() self.feature_extractor = nn.Sequential( AnalogConv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, rpu_config=RPU_CONFIG), nn.Tanh(), nn.MaxPool2d(kernel_size=2), AnalogConv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, rpu_config=RPU_CONFIG), nn.Tanh(), nn.MaxPool2d(kernel_size=2), nn.Tanh()) self.classifier = nn.Sequential( AnalogLinear(in_features=512, out_features=128, rpu_config=RPU_CONFIG), nn.Tanh(), AnalogLinear(in_features=128, out_features=N_CLASSES, rpu_config=RPU_CONFIG), )
def get_model(self, rpu_config: Any = TikiTakaEcRamPreset) -> Module: return AnalogSequential( Conv2d(in_channels=3, out_channels=48, kernel_size=3, stride=1, padding=1), ReLU(), AnalogConv2d(in_channels=48, out_channels=48, kernel_size=3, stride=1, padding=1, rpu_config=rpu_config(), weight_scaling_omega=0.8), BatchNorm2d(48), ReLU(), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=48, out_channels=96, kernel_size=3, stride=1, padding=1, rpu_config=rpu_config(), weight_scaling_omega=0.8), ReLU(), AnalogConv2d(in_channels=96, out_channels=96, kernel_size=3, stride=1, padding=1, rpu_config=rpu_config(), weight_scaling_omega=0.8), BatchNorm2d(96), ReLU(), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=96, out_channels=144, kernel_size=3, stride=1, padding=1, rpu_config=rpu_config(), weight_scaling_omega=0.8), ReLU(), AnalogConv2d(in_channels=144, out_channels=144, kernel_size=3, stride=1, padding=1, rpu_config=rpu_config(), weight_scaling_omega=0.8), BatchNorm2d(144), ReLU(), MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), Flatten(), AnalogLinear(in_features=16 * 144, out_features=384, rpu_config=rpu_config(), weight_scaling_omega=0.8), ReLU(), Linear(in_features=384, out_features=10), LogSoftmax(dim=1))
def get_layer(self, in_channels=2, out_channels=3, **kwargs): kwargs.setdefault('rpu_config', self.get_rpu_config()) kwargs.setdefault('bias', self.bias) kwargs.setdefault('kernel_size', [3, 3]) kwargs.setdefault('padding', 2) kwargs['rpu_config'].mapping.digital_bias = self.digital_bias return AnalogConv2d(in_channels, out_channels, **kwargs)
def get_layer(self, in_channels=2, out_channels=3, kernel_size=4, padding=2, **kwargs): kwargs.setdefault('rpu_config', self.get_rpu_config()) kwargs.setdefault('bias', self.bias) return AnalogConv2d(in_channels, out_channels, kernel_size, padding=padding, **kwargs).cuda()
def get_model(self, rpu_config: Any = ReRamSBPreset) -> Module: return AnalogSequential( AnalogConv2d(in_channels=3, out_channels=16, kernel_size=5, stride=1, rpu_config=rpu_config()), Tanh(), MaxPool2d(kernel_size=2), AnalogConv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, rpu_config=rpu_config()), Tanh(), MaxPool2d(kernel_size=2), Tanh(), Flatten(), AnalogLinear(in_features=800, out_features=128, rpu_config=rpu_config()), Tanh(), AnalogLinear(in_features=128, out_features=10, rpu_config=rpu_config()), LogSoftmax(dim=1))
def create_analog_network(): """Return a LeNet5 inspired analog model.""" channel = [16, 32, 512, 128] model = AnalogSequential( AnalogConv2d(in_channels=1, out_channels=channel[0], kernel_size=5, stride=1, rpu_config=RPU_CONFIG), nn.Tanh(), nn.MaxPool2d(kernel_size=2), AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=5, stride=1, rpu_config=RPU_CONFIG), nn.Tanh(), nn.MaxPool2d(kernel_size=2), nn.Tanh(), nn.Flatten(), AnalogLinear(in_features=channel[2], out_features=channel[3], rpu_config=RPU_CONFIG), nn.Tanh(), AnalogLinear(in_features=channel[3], out_features=N_CLASSES, rpu_config=RPU_CONFIG), nn.LogSoftmax(dim=1) ) return model
def get_model(self, rpu_config: Any = EcRamPreset) -> Module: return AnalogSequential( AnalogConv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, rpu_config=rpu_config(), weight_scaling_omega=0.6), Tanh(), MaxPool2d(kernel_size=2), AnalogConv2d(in_channels=16, out_channels=32, kernel_size=5, stride=1, rpu_config=rpu_config(), weight_scaling_omega=0.6), Tanh(), MaxPool2d(kernel_size=2), Tanh(), Flatten(), AnalogLinear(in_features=512, out_features=128, rpu_config=rpu_config(), weight_scaling_omega=0.6), Tanh(), AnalogLinear(in_features=128, out_features=10, rpu_config=rpu_config(), weight_scaling_omega=0.6), LogSoftmax(dim=1))
from torch import nn # Imports from aihwkit. from aihwkit.nn import AnalogLinear, AnalogConv2d, AnalogSequential from aihwkit.simulator.configs import SingleRPUConfig from aihwkit.simulator.configs.devices import ConstantStepDevice from aihwkit.utils.analog_info import analog_summary # Define a single-layer network, using a constant step device type. rpu_config = SingleRPUConfig(device=ConstantStepDevice()) channel = [16, 32, 512, 128] model = AnalogSequential( AnalogConv2d(in_channels=1, out_channels=channel[0], kernel_size=5, stride=1, rpu_config=rpu_config), nn.Tanh(), nn.MaxPool2d(kernel_size=2), AnalogConv2d(in_channels=channel[0], out_channels=channel[1], kernel_size=5, stride=1, rpu_config=rpu_config), nn.Tanh(), nn.MaxPool2d(kernel_size=2), nn.Tanh(), nn.Flatten(), AnalogLinear(in_features=channel[2], out_features=channel[3], rpu_config=rpu_config), nn.Tanh(), AnalogLinear(in_features=channel[3], out_features=10, rpu_config=rpu_config), nn.LogSoftmax(dim=1))
def VGG8(): """VGG8 inspired analog model.""" model = AnalogSequential( AnalogConv2d(in_channels=3, out_channels=128, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=128, out_channels=128, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(128), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=128, out_channels=256, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=256, out_channels=256, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(256), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), AnalogConv2d(in_channels=256, out_channels=512, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogConv2d(in_channels=512, out_channels=512, kernel_size=3, stride=1, padding=1, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.BatchNorm2d(512), nn.ReLU(), nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1), nn.Flatten(), AnalogLinear(in_features=8192, out_features=1024, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.ReLU(), AnalogLinear(in_features=1024, out_features=N_CLASSES, rpu_config=RPU_CONFIG, weight_scaling_omega=WEIGHT_SCALING_OMEGA), nn.LogSoftmax(dim=1)) return model