torch.set_num_threads(os.cpu_count() - 1)
print("Running on Device = ", device)
# Determines number of workers to use
if n_workers == -1:
n_workers = torch.cuda.is_available() * 4 * torch.cuda.device_count()
n_sqrt = int(np.ceil(np.sqrt(n_neurons)))
start_intensity = intensity
# Build network.
network = IncreasingInhibitionNetwork(
n_input=784,
n_neurons=n_neurons,
start_inhib=10,
max_inhib=-40.0,
theta_plus=0.05,
tc_theta_decay=1e7,
inpt_shape=(1, 28, 28),
nu=(1e-4, 1e-2),
)
network.to(device)
# Load MNIST data.
dataset = MNIST(
PoissonEncoder(time=time, dt=dt),
None,
root=os.path.join("..", "..", "data", "MNIST"),
download=True,
transform=transforms.Compose(
[transforms.ToTensor(),
device = 'cuda' if gpu else 'cpu'
n_examples = n_train if train else n_test
n_sqrt = int(np.ceil(np.sqrt(n_neurons)))
start_intensity = intensity
n_classes = 4
per_class = int(n_examples / n_classes)
n_low = int(p_low * n_train)
# Build network.
if train:
network = IncreasingInhibitionNetwork(n_input=50 * 72,
n_neurons=n_neurons,
start_inhib=start_inhib,
max_inhib=max_inhib,
dt=dt,
norm=norm,
theta_plus=theta_plus,
theta_decay=theta_decay)
else:
network = load_network(os.path.join(params_path, model_name + '.pt'))
network.connections[('X', 'Y')].update_rule = NoOp(
connection=network.connections[('X', 'Y')])
# Load Breakout data.
images = torch.load(os.path.join(data_path, 'frames.pt'),
map_location=torch.device(device))
labels = torch.load(os.path.join(data_path, 'labels.pt'),
map_location=torch.device(device))
images = images[:, 30:, 4:-4].contiguous().view(
-1, 50 * 72) # Crop out the borders of the frames.