def test_apply_noise(self):
"""Test using realistic weights (bias)."""
weights = randn(10, 35)
noise_model = PCMLikeNoiseModel()
t_inference = 100.
noisy_weights = noise_model.apply_noise(weights, t_inference)
self.assertNotAlmostEqualTensor(noisy_weights, weights)
def get_model_and_x(self):
"""Trains a simple model."""
# Prepare the datasets (input and expected output).
x = Tensor([[0.1, 0.2, 0.4, 0.3], [0.2, 0.1, 0.1, 0.3]])
y = Tensor([[1.0, 0.5], [0.7, 0.3]])
# Define a single-layer network, using a constant step device type.
rpu_config = self.get_rpu_config()
rpu_config.forward.out_res = -1. # Turn off (output) ADC discretization.
rpu_config.forward.w_noise_type = WeightNoiseType.ADDITIVE_CONSTANT
rpu_config.forward.w_noise = 0.02
rpu_config.noise_model = PCMLikeNoiseModel(g_max=25.0)
model = AnalogLinear(4, 2, bias=True, rpu_config=rpu_config)
# Move the model and tensors to cuda if it is available.
if self.use_cuda:
x = x.cuda()
y = y.cuda()
model.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)
opt.regroup_param_groups(model)
for _ in range(100):
# Add the training Tensor to the model (input).
pred = model(x)
# Add the expected output Tensor.
loss = mse_loss(pred, y)
# Run training (backward propagation).
loss.backward()
opt.step()
return model, x
RESULTS = os.path.join('results', 'LENET5')
# Training parameters
SEED = 1
N_EPOCHS = 30
BATCH_SIZE = 8
LEARNING_RATE = 0.01
N_CLASSES = 10
# Define the properties of the neural network in terms of noise simulated during
# the inference/training pass
RPU_CONFIG = InferenceRPUConfig()
RPU_CONFIG.forward.out_res = -1. # Turn off (output) ADC discretization.
RPU_CONFIG.forward.w_noise_type = WeightNoiseType.ADDITIVE_CONSTANT
RPU_CONFIG.forward.w_noise = 0.02
RPU_CONFIG.noise_model = PCMLikeNoiseModel(g_max=25.0)
def load_images():
"""Load images for train from torchvision datasets."""
transform = transforms.Compose([transforms.ToTensor()])
train_set = datasets.MNIST(TRAIN_DATASET, download=True, train=True, transform=transform)
val_set = datasets.MNIST(TEST_DATASET, download=True, train=False, transform=transform)
train_data = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=True)
validation_data = torch.utils.data.DataLoader(val_set, batch_size=BATCH_SIZE, shuffle=False)
return train_data, validation_data
class LeNet5(AnalogSequential):
"""LeNet5 inspired analog model."""
# Define a single-layer network, using inference/hardware-aware training tile
rpu_config = InferenceRPUConfig()
rpu_config.forward.out_res = -1. # Turn off (output) ADC discretization.
rpu_config.forward.w_noise_type = OutputWeightNoiseType.ADDITIVE_CONSTANT
rpu_config.forward.w_noise = 0.02 # Short-term w-noise.
rpu_config.clip.type = WeightClipType.FIXED_VALUE
rpu_config.clip.fixed_value = 1.0
rpu_config.modifier.pdrop = 0.03 # Drop connect.
rpu_config.modifier.type = WeightModifierType.ADD_NORMAL # Fwd/bwd weight noise.
rpu_config.modifier.std_dev = 0.1
rpu_config.modifier.rel_to_actual_wmax = True
# Inference noise model.
rpu_config.noise_model = PCMLikeNoiseModel(g_max=25.0)
# drift compensation
rpu_config.drift_compensation = GlobalDriftCompensation()
model = AnalogLinear(4, 2, bias=True,
rpu_config=rpu_config)
# Move the model and tensors to cuda if it is available.
if cuda.is_compiled():
x = x.cuda()
y = y.cuda()
model.cuda()
# Define an analog-aware optimizer, preparing it for using the layers.
opt = AnalogSGD(model.parameters(), lr=0.1)