Commit f2bffc72 authored by Devanshu Agrawal's avatar Devanshu Agrawal
Browse files

Make no-bottleneck NNGP work.

parent 12b103e9

torch_experiments/main.py

+13 −2
Original line number Diff line number Diff line
@@ -17,7 +17,7 @@ parser=argparse.ArgumentParser() 
parser.add_argument("--dataset", "-ds", required=True, help="Dataset. One of: boston, iris, rings.")

# architecture

parser.add_argument("--depths", "-d", type=lambda s: [int(i) for i in s.split(",")], default="0", help="List of depths (number of hidden layers) of the NNGP components in comma-separated format.")

parser.add_argument("--widths", "-w", type=lambda s: [int(i) for i in s.split(",")], default="", help="List of bottleneck widths in comma-separated format.")

parser.add_argument("--widths", "-w", type=lambda s: [int(i) for i in s.split(",")], default="", help="List of bottleneck widths in comma-separated format. A width of 0 will be interpreted as infinity.")

# initial variance hyperparameters

parser.add_argument('--vb','-vb', default=1.0, type=float, help="Initial value of hyperparameter v_b.")

parser.add_argument('--vw','-vw', default=1.0, type=float, help="Initial value of hyperparameter v_w.")
@@ -35,6 +35,17 @@ parser.add_argument("--gpu", "-g", default=-1, type=int, help="Which GPU to use. 


args = parser.parse_args()

device = "cpu" if args.gpu < 0 else "cuda:{}".format(args.gpu)

# get depths and widths adjusting for width 0 (i.e., infinite) layers

depths, widths = [], []

depth = args.depths[0]

for (d, w) in zip(args.depths[1:], args.widths):

	if w == 0:

		depth += d+1

	else:

		depths.append(depth)

		widths.append(w)

		depth = d

depths.append(depth)





# load dataset
@@ -43,7 +54,7 @@ dataloader = DataLoader(dataset, batch_size=len(dataset)) 
X, Y = next(iter(dataloader))



# get model

model = BottleneckNNGP(depths=args.depths, widths=args.widths, v_b=args.vb, v_w=args.vw, v_n=args.vn, device=device, manual_grad=args.manual_grad, to_dtype=torch.float32, to_device="cpu")

model = BottleneckNNGP(depths=depths, widths=widths, v_b=args.vb, v_w=args.vw, v_n=args.vn, device=device, manual_grad=args.manual_grad, to_dtype=torch.float32, to_device="cpu")



# define stochastic loss function

loss_fn = lambda num_samples, manual_samples: -model.log_likelihood(X, Y, num_samples=num_samples, manual_samples=manual_samples)/X.shape[0]

torch_experiments/modules.py

+1 −1
Original line number Diff line number Diff line
@@ -71,7 +71,7 @@ class BottleneckNNGP(nn.Module): 
			samples = torch.tensor([2.], dtype=self.dtype, device=self.device).sqrt() * nn.functional.relu(samples)

			gram_matrices = torch.matmul(samples, torch.transpose(samples, 1, 2)) / width



		if not torch.is_grad_enabled():

		if not torch.is_grad_enabled() and len(self.widths) > 0:

			del Ks

			del Ls

			del samples