def _create_input(self):
"""Return same random input (reset seed before) and labels."""
set_seeds(self.SEED)
x = torch.randn(*self.INPUT_SIZE).to(self.DEVICE)
y = torch.randint(size=(self.BATCH, ),
low=0,
high=self.NUM_CLASSES).to(self.DEVICE)
return x, y
def create_layers():
"""Create list of layers for the example sequence."""
set_seeds(0)
layers = []
for cls, in_, out in zip(classes, in_features, out_features):
layers.append(cls(in_features=in_, out_features=out, bias=True))
for cls in classes[3:]:
layers.append(cls())
return layers
def create_layers(activation):
"""Create layers of the fully-connected network."""
# same seed
set_seeds(0)
layers = []
for (in_, out) in zip(in_features, out_features):
layers.append(HBPLinear(in_features=in_, out_features=out))
layers.append(activation())
return layers
def torch_fn():
"""Create a 2d convolution layer in torch."""
set_seeds(0)
return Conv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=bias,
)
def cvp_fn():
"""Create a 2d convolution layer with CVP functionality."""
set_seeds(0)
return CVPConv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=bias,
)
def test_deterministic_loading():
"""Test deterministc loading of samples."""
set_seeds(13)
cifar10_1 = CIFAR10Loader(10, 10)
train1 = cifar10_1.train_loader()
train1_samples, train1_labels = next(iter(train1))
set_seeds(13)
cifar10_2 = CIFAR10Loader(10, 10)
train2 = cifar10_2.train_loader()
train2_samples, train2_labels = next(iter(train2))
assert allclose(train1_samples, train2_samples)
assert allclose(train1_labels, train2_labels)
def test_deterministic_loading():
"""Test deterministc loading of samples."""
for loader in loaders:
set_seeds(0)
mnist1 = loader(10, 10)
train1 = mnist1.train_loader()
train1_samples, train1_labels = next(iter(train1))
set_seeds(0)
mnist2 = loader(10, 10)
train2 = mnist2.train_loader()
train2_samples, train2_labels = next(iter(train2))
assert allclose(train1_samples, train2_samples)
assert allclose(train1_labels, train2_labels)
def hbp_fn():
"""Create sequence of layers in HBP."""
set_seeds(0)
return HBPSequential(
HBPConv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=True,
),
HBPReLU(),
HBPMaxPool2d(pool_kernel, padding=pool_padding),
HBPFlatten(),
HBPLinear(out1, out2, bias=False),
HBPSigmoid(),
HBPLinear(out2, out3, bias=True),
)
def torch_fn():
"""Create sequence of layers in torch."""
set_seeds(0)
return Sequential(
Conv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=True,
),
ReLU(),
MaxPool2d(pool_kernel, padding=pool_padding),
Flatten(),
Linear(out1, out2, bias=False),
Sigmoid(),
Linear(out2, out3, bias=True),
)
def training_example(seed, test_batch, use_gpu=False):
"""Training instance setting seed and test batch size in advance."""
set_seeds(seed)
device = torch.device("cuda:0" if use_gpu else "cpu")
model = Sequential(Flatten(), Linear(784, 10))
loss_function = CrossEntropyLoss()
data_loader = MNISTLoader(1000, test_batch)
optimizer = SGD(model.parameters(), lr=0.1)
# initialize training
train = FirstOrderTraining(
model,
loss_function,
optimizer,
data_loader,
logdir,
num_epochs,
logs_per_epoch=logs_per_epoch,
device=device,
)
return train
def cvp_fn():
"""Create sequence of layers in torch."""
set_seeds(0)
return CVPSequential(
CVPConv2d(
in_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
bias=True,
),
CVPReLU(),
CVPMaxPool2d(kernel_size,
stride=stride,
padding=padding,
dilation=dilation),
CVPFlatten(),
CVPLinear(out2_features, out3_features, bias=False),
CVPSigmoid(),
CVPLinear(out3_features, out4_features, bias=True),
)
def test_forward_cifar10_models():
"""Check same behaviour of original and HBP/split CIFAR-10 model."""
max_blocks = 5
input = torch.randn(2, 3072)
set_seeds(0)
original = cifar10_model()
set_seeds(0)
hbp = hbp_cifar10_model()
set_seeds(0)
hbp_parallel = hbp_split_cifar10_model(max_blocks, False, False)
assert torch.allclose(original(input), hbp(input), atol=1e-5)
assert torch.allclose(original(input), hbp_parallel(input), atol=1e-5)
def test_forward_mnist_models():
"""Check same behaviour of original and HBP/split MNIST model."""
max_blocks = 5
input = torch.randn(2, 784)
set_seeds(0)
original = mnist_model()
set_seeds(0)
hbp = hbp_mnist_model()
set_seeds(0)
hbp_parallel = hbp_split_mnist_model(max_blocks)
assert torch.allclose(original(input), hbp(input), atol=1e-5)
assert torch.allclose(original(input), hbp_parallel(input), atol=1e-5)
def _create_input(self):
"""Return same random input (reset seed before)."""
set_seeds(self.SEED)
return torch.randn(*self.INPUT_SIZE).to(self.DEVICE)
import matplotlib.pyplot as plt
import torch
import torchvision
import torchvision.transforms as transforms
from bpexts.hbp.crossentropy import HBPCrossEntropyLoss
from bpexts.hbp.linear import HBPLinear
from bpexts.hbp.sequential import HBPSequential
from bpexts.hbp.sigmoid import HBPSigmoid
from bpexts.optim.cg_newton import CGNewton
from bpexts.utils import set_seeds
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
set_seeds(0)
batch_size = 500
# download directory
data_dir = "~/tmp/CIFAR10"
# training set loader
train_set = torchvision.datasets.CIFAR10(root=data_dir,
train=True,
transform=transforms.ToTensor(),
download=True)
train_loader = torch.utils.data.DataLoader(dataset=train_set,
batch_size=batch_size,
shuffle=True)
# layers
def hbp_fn():
"""Create a linear layer with HBP functionality."""
set_seeds(0)
return HBPLinear(in_features=in_features,
out_features=out_features,
bias=bias)
def torch_fn():
"""Create a linear layer in torch."""
set_seeds(0)
return Linear(in_features=in_features,
out_features=out_features,
bias=bias)
def example_linear():
"""Return example layer of HBPLinear."""
set_seeds(0)
return HBPLinear(in_features=in_features,
out_features=out_features,
bias=True)
