def example_layer():
"""Return example linear layer."""
layer = HBPLinear(in_features=3, out_features=2)
w = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
requires_grad=True)
b = torch.tensor([7.0, 8.0], requires_grad=True)
layer.weight.data = w
layer.bias.data = b
return layer
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 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 simple_mnist_model_2nd_order_hbp(use_gpu=False):
"""Train on simple MNIST model using 2nd order optimizer HBP."""
device = torch.device("cuda:0" if use_gpu else "cpu")
model = HBPSequential(HBPFlatten(), HBPLinear(784, 10))
loss_function = convert_torch_to_hbp(CrossEntropyLoss())
data_loader = MNISTLoader(1000, 1000)
optimizer = CGNewton(model.parameters(), lr=0.1, alpha=0.1)
num_epochs, logs_per_epoch = 1, 5
modify_2nd_order_terms = "abs"
logdir = directory_in_data("test_training_simple_mnist_model")
# initialize training
train = HBPSecondOrderTraining(
model,
loss_function,
optimizer,
data_loader,
logdir,
num_epochs,
modify_2nd_order_terms,
logs_per_epoch=logs_per_epoch,
device=device,
)
train.run()
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
linear1 = HBPLinear(in_features=3072, out_features=1024, bias=True)
activation1 = HBPSigmoid()
linear2 = HBPLinear(in_features=1024, out_features=512, bias=True)
activation2 = HBPSigmoid()
linear3 = HBPLinear(in_features=512, out_features=256, bias=True)
activation3 = HBPSigmoid()
linear4 = HBPLinear(in_features=256, out_features=128, bias=True)
activation4 = HBPSigmoid()
linear5 = HBPLinear(in_features=128, out_features=10, bias=True)
# sequential model
model = HBPSequential(
linear1,
activation1,
linear2,
activation2,
def hbp_from_torch_fn():
"""Create HBLinear from Linear."""
torch_layer = torch_fn()
return HBPLinear.from_torch(torch_layer)
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)
train_set = torchvision.datasets.MNIST(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)
# layer parameters
in_features = 784
out_features = 10
bias = True
# linear layer
model = HBPLinear(in_features=in_features,
out_features=out_features,
bias=bias)
# load to device
model.to(device)
print(model)
loss_func = HBPCrossEntropyLoss()
# learning rate
lr = 0.1
# regularization
alpha = 0.02
# convergence criteria for CG
cg_maxiter = 50
def example_linear():
"""Return example layer of HBPLinear."""
set_seeds(0)
return HBPLinear(in_features=in_features,
out_features=out_features,
bias=True)