for data, target in train_loader2:
data = data.to(device)
target = target.to(device)
output = model(data)
pred = output.data.max(1)[1] # get the index of the max log-probability
correct += pred.eq(target.data).cpu().sum()
ANN_accuracy = 100. * correct.to(torch.float32) / len(train_loader2.dataset)
print("ANN accuracy:", ANN_accuracy)
validate()
start = t_()
for index, (data, target) in enumerate(train_loader2):
print('sample ', index+1, 'elapsed', t_() - start)
start = t_()
data = data.to(device)
data = data.view(-1, 28*28)
inpts = {'Input': data.repeat(time, 1)}
SNN.run(inpts=inpts, time=time)
spikes = {layer: SNN.monitors[layer].get('s') for layer in SNN.monitors}
voltages = {layer: SNN.monitors[layer].get('v') for layer in SNN.monitors if not layer == 'Input'}
pred = torch.argmax(voltages['2'].sum(1))
correct += pred.eq(target.data.to(device)).cpu().sum()
accuracy = 100. * correct.to(torch.float32) / (index + 1)
SNN.reset_()
from time import time as t_
t0 = t_()
def gs(pix):
return pix[0] * 30 + pix[1] * 59 + pix[2] * 11
red, teal, navy, offwhite, green, deepred, banana, deepblue, white = (
218, 20, 21), (112, 150,
160), (0, 48, 80), (250, 227,
173), (176, 183,
167), (69, 0,
0), (69 + 69 + 69, 169,
69), (0, 34,
69), (255, 255,
255)
g_ = [(0, 255, 0), (0, 64, 32), (24, 128, 48), (160, 255, 192),
(200, 255, 200)]
r_ = [(i[1], i[2], i[0]) for i in g_]
b_ = [(i[2], i[0], i[1]) for i in g_]
b_ = [(255, 0, 255), (96, 0, 96), (192, 36, 192), (255, 180, 255),
(255, 210, 255)]
_ = [(sum(i) // 3, sum(i) // 3, sum(i) // 3) for i in g_]
presets = {
"original": [red, teal, navy, offwhite, green],
"sun": [red, offwhite, deepred, banana],
"dark": [teal, green, offwhite, deepblue, white],
def main(seed=0, time=50, n_episodes=25, n_snn_episodes=100, percentile=99.9, epsilon=0.05, occlusion=0, plot=False):
np.random.seed(seed)
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
print()
print('Loading the trained ANN...')
print()
ANN = Net()
ANN.load_state_dict(
torch.load(
'../../params/pytorch_breakout_dqn.pt'
)
)
environment = make_atari('BreakoutNoFrameskip-v4')
environment = wrap_deepmind(environment, frame_stack=True, scale=False, clip_rewards=False, episode_life=False)
f = f'{seed}_{n_episodes}_states.pt'
if os.path.isfile(os.path.join(params_path, f)):
print('Loading pre-gathered observation data...')
states = torch.load(os.path.join(params_path, f))
else:
print('Gathering observation data...')
print()
episode_rewards = np.zeros(n_episodes)
total_t = 0
states = []
for i in range(n_episodes):
state = torch.tensor(environment.reset()).to(device).unsqueeze(0).permute(0, 3, 1, 2).float()
for t in itertools.count():
states.append(state)
q_values = ANN(state)[0]
probs, best_action = policy(q_values, epsilon)
action = np.random.choice(np.arange(len(probs)), p=probs)
state, reward, done, _ = environment.step(action)
state = torch.tensor(state).unsqueeze(0).permute(0, 3, 1, 2).float()
state = state.to(device)
episode_rewards[i] += reward
total_t += 1
if done:
print(f'Step {t} ({total_t}) @ Episode {i + 1} / {n_episodes}')
print(f'Episode Reward: {episode_rewards[i]}')
break
states = torch.cat(states, dim=0)
torch.save(states, os.path.join(params_path, f))
print()
print(f'Collected {states.size(0)} Atari game frames.')
print()
print('Converting ANN to SNN...')
states = states.to(device)
# Do ANN to SNN conversion.
SNN = ann_to_snn(ANN, input_shape=(1, 4, 84, 84), data=states / 255.0, percentile=percentile)
for l in SNN.layers:
if l != 'Input':
SNN.add_monitor(
Monitor(SNN.layers[l], state_vars=['s', 'v'], time=time), name=l
)
else:
SNN.add_monitor(
Monitor(SNN.layers[l], state_vars=['s'], time=time), name=l
)
spike_ims = None
spike_axes = None
inpt_ims = None
inpt_axes = None
voltage_ims = None
voltage_axes = None
new_life = True
rewards = np.zeros(n_snn_episodes)
total_t = 0
noop_counter = 0
print()
print('Testing SNN on Atari Breakout game...')
print()
# Test SNN on Atari Breakout.
for i in range(n_snn_episodes):
state = torch.tensor(environment.reset()).to(device).unsqueeze(0).permute(0, 3, 1, 2)
prev_life = 5
start = t_()
for t in itertools.count():
print(f'Timestep {t} (elapsed {t_() - start:.2f})')
start = t_()
sys.stdout.flush()
state[:, :, 77 - occlusion: 80 - occlusion, :] = 0
import matplotlib.pyplot as plt
print(state.size())
plt.matshow(state.float().mean(1).squeeze(0).cpu())
plt.ioff()
plt.show()
state = state.repeat(time, 1, 1, 1, 1)
inpts = {'Input': state.float() / 255.0}
SNN.run(inpts=inpts, time=time)
spikes = {layer: SNN.monitors[layer].get('s') for layer in SNN.monitors}
voltages = {layer: SNN.monitors[layer].get('v') for layer in SNN.monitors if not layer == 'Input'}
probs, best_action = policy(voltages['12'].sum(1), epsilon)
action = np.random.choice(np.arange(len(probs)), p=probs)
if action == 0:
noop_counter += 1
else:
noop_counter = 0
if noop_counter >= 20:
action = np.random.choice([0, 1, 2, 3])
noop_counter = 0
if new_life:
action = 1
next_state, reward, done, info = environment.step(action)
next_state = torch.tensor(next_state).unsqueeze(0).permute(0, 3, 1, 2)
if prev_life - info["ale.lives"] != 0:
new_life = True
else:
new_life = False
prev_life = info["ale.lives"]
rewards[i] += reward
total_t += 1
SNN.reset_()
if plot:
# Get voltage recording.
inpt = state.view(time, 4, 84, 84).sum(0).sum(0).view(84, 84)
spike_ims, spike_axes = plot_spikes(
{layer: spikes[layer] for layer in spikes}, ims=spike_ims, axes=spike_axes
)
voltage_ims, voltage_axes = plot_voltages(
{layer: voltages[layer].view(time, -1) for layer in voltages},
ims=voltage_ims, axes=voltage_axes
)
inpt_axes, inpt_ims = plot_input(inpt, inpt, ims=inpt_ims, axes=inpt_axes)
plt.pause(1e-8)
if done:
print(f'Step {t} ({total_t}) @ Episode {i + 1} / {n_snn_episodes}')
print(f'Episode Reward: {rewards[i]}')
print()
break
state = next_state
model_name = '_'.join([str(x) for x in [seed, time, n_episodes, n_snn_episodes, percentile, epsilon, occlusion]])
columns = [
'seed', 'time', 'n_episodes', 'n_snn_episodes', 'percentile', 'epsilon', 'occlusion', 'avg. reward', 'std. reward'
]
data = [[
seed, time, n_episodes, n_snn_episodes, percentile, epsilon, occlusion, np.mean(rewards), np.std(rewards)
]]
path = os.path.join(results_path, 'results.csv')
if not os.path.isfile(path):
df = pd.DataFrame(data=data, index=[model_name], columns=columns)
else:
df = pd.read_csv(path, index_col=0)
if model_name not in df.index:
df = df.append(pd.DataFrame(data=data, index=[model_name], columns=columns))
else:
df.loc[model_name] = data[0]
df.to_csv(path, index=True)
torch.save(rewards, os.path.join(results_path, f'{model_name}_episode_rewards.pt'))
def main(occlusion_percentage=0):
percentile = 99.9
random_seed = 0
torch.manual_seed(random_seed)
time = 100
class RandomlyOcclude(object):
def __init__(self, percentage):
self.percentage = percentage / 100
def __call__(self, img):
mask = torch.Tensor(img.shape[1],
img.shape[2]).uniform_() > self.percentage
return img.to(device) * mask.float()
train_dataset2 = datasets.MNIST('./data',
train=False,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
RandomlyOcclude(occlusion_percentage)
]))
train_dataset = datasets.MNIST('./data',
train=True,
download=True,
transform=transforms.ToTensor())
train_loader2 = torch.utils.data.DataLoader(dataset=train_dataset2,
shuffle=True)
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
shuffle=True,
batch_size=train_dataset.__len__())
for d, target in train_loader:
data = d.to(device)
model = torch.load('trained_model.pt')
print()
print('Converting ANN to SNN...')
SNN = ann_to_snn(model,
input_shape=[28 * 28],
data=data,
percentile=percentile)
SNN.add_monitor(Monitor(SNN.layers['2'], state_vars=['s', 'v'], time=time),
name='2')
correct = 0
def validate():
model.eval()
val_loss, correct = 0, 0
for data, target in train_loader2:
data = data.to(device)
target = target.to(device)
output = model(data)
pred = output.data.max(1)[
1] # get the index of the max log-probability
correct += pred.eq(target.data).cpu().sum()
ANN_accuracy = 100. * correct.to(torch.float32) / len(
train_loader2.dataset)
print("ANN accuracy:", ANN_accuracy)
return ANN_accuracy
ANN_accuracy = validate()
start = t_()
for index, (data, target) in enumerate(train_loader2):
print('sample ', index + 1, 'elapsed', t_() - start)
start = t_()
data = data.to(device)
data = data.view(-1, 28 * 28)
inpts = {'Input': data.repeat(time, 1)}
SNN.run(inpts=inpts, time=time)
spikes = {
layer: SNN.monitors[layer].get('s')
for layer in SNN.monitors
}
voltages = {
layer: SNN.monitors[layer].get('v')
for layer in SNN.monitors if not layer == 'Input'
}
pred = torch.argmax(voltages['2'].sum(1))
correct += pred.eq(target.data.to(device)).cpu().sum()
SNN.reset_()
SNN_accuracy = 100. * correct.to(torch.float32) / len(
train_loader2.dataset)
print("accuracy:, ", SNN_accuracy)
df = pd.DataFrame({
"ANN accuracy": [ANN_accuracy],
"SNN accuracy": [SNN_accuracy]
})
df.to_csv("accuracy_1hidden_" + str(occlusion_percentage) + ".csv")
def main(seed=0,
time=250,
n_snn_episodes=1,
epsilon=0.05,
plot=False,
parameter1=1.0,
parameter2=1.0,
parameter3=1.0,
parameter4=1.0,
parameter5=1.0):
np.random.seed(seed)
parameters = [parameter1, parameter2, parameter3, parameter4, parameter5]
if torch.cuda.is_available():
torch.set_default_tensor_type('torch.cuda.FloatTensor')
torch.cuda.manual_seed_all(seed)
else:
torch.manual_seed(seed)
print()
print('Loading the trained ANN...')
print()
ANN = Net()
ANN.load_state_dict(torch.load('../../params/pytorch_breakout_dqn.pt'))
environment = make_atari('BreakoutNoFrameskip-v4')
environment = wrap_deepmind(environment,
frame_stack=True,
scale=False,
clip_rewards=False,
episode_life=False)
print('Converting ANN to SNN...')
# Do ANN to SNN conversion.
# SNN = ann_to_snn(ANN, input_shape=(1, 4, 84, 84), data=states / 255.0, percentile=percentile, node_type=LIFNodes, decay=1e-2 / 13.0, rest=0.0)
SNN = Network()
input_layer = nodes.RealInput(shape=(1, 4, 84, 84))
SNN.add_layer(input_layer, name='Input')
children = []
for c in ANN.children():
if isinstance(c, nn.Sequential):
for c2 in list(c.children()):
children.append(c2)
else:
children.append(c)
i = 0
prev = input_layer
scale_index = 0
while i < len(children) - 1:
current, nxt = children[i:i + 2]
layer, connection = _ann_to_snn_helper(prev,
current,
scale=parameters[scale_index])
i += 1
if layer is None or connection is None:
continue
SNN.add_layer(layer, name=str(i))
SNN.add_connection(connection, source=str(i - 1), target=str(i))
prev = layer
if isinstance(current, nn.Linear) or isinstance(current, nn.Conv2d):
scale_index += 1
current = children[-1]
layer, connection = _ann_to_snn_helper(prev,
current,
scale=parameters[scale_index])
i += 1
if layer is not None or connection is not None:
SNN.add_layer(layer, name=str(i))
SNN.add_connection(connection, source=str(i - 1), target=str(i))
for l in SNN.layers:
if l != 'Input':
SNN.add_monitor(Monitor(SNN.layers[l],
state_vars=['s', 'v'],
time=time),
name=l)
else:
SNN.add_monitor(Monitor(SNN.layers[l], state_vars=['s'],
time=time),
name=l)
spike_ims = None
spike_axes = None
inpt_ims = None
inpt_axes = None
voltage_ims = None
voltage_axes = None
rewards = np.zeros(n_snn_episodes)
total_t = 0
print()
print('Testing SNN on Atari Breakout game...')
print()
# Test SNN on Atari Breakout.
for i in range(n_snn_episodes):
state = torch.tensor(
environment.reset()).to(device).unsqueeze(0).permute(0, 3, 1, 2)
start = t_()
for t in itertools.count():
print(f'Timestep {t} (elapsed {t_() - start:.2f})')
start = t_()
sys.stdout.flush()
state = state.repeat(time, 1, 1, 1, 1)
inpts = {'Input': state.float() / 255.0}
SNN.run(inpts=inpts, time=time)
spikes = {
layer: SNN.monitors[layer].get('s')
for layer in SNN.monitors
}
voltages = {
layer: SNN.monitors[layer].get('v')
for layer in SNN.monitors if not layer == 'Input'
}
probs, best_action = policy(spikes['12'].sum(1), epsilon)
action = np.random.choice(np.arange(len(probs)), p=probs)
next_state, reward, done, info = environment.step(action)
next_state = torch.tensor(next_state).unsqueeze(0).permute(
0, 3, 1, 2)
rewards[i] += reward
total_t += 1
SNN.reset_()
if plot:
# Get voltage recording.
inpt = state.view(time, 4, 84, 84).sum(0).sum(0).view(84, 84)
spike_ims, spike_axes = plot_spikes(
{layer: spikes[layer]
for layer in spikes},
ims=spike_ims,
axes=spike_axes)
voltage_ims, voltage_axes = plot_voltages(
{
layer: voltages[layer].view(time, -1)
for layer in voltages
},
ims=voltage_ims,
axes=voltage_axes)
inpt_axes, inpt_ims = plot_input(inpt,
inpt,
ims=inpt_ims,
axes=inpt_axes)
plt.pause(1e-8)
if done:
print(
f'Step {t} ({total_t}) @ Episode {i + 1} / {n_snn_episodes}'
)
print(f'Episode Reward: {rewards[i]}')
print()
break
state = next_state
model_name = '_'.join([
str(x) for x in
[seed, parameter1, parameter2, parameter3, parameter4, parameter5]
])
columns = [
'seed', 'time', 'n_snn_episodes', 'avg. reward', 'parameter1',
'parameter2', 'parameter3', 'parameter4', 'parameter5'
]
data = [[
seed, time, n_snn_episodes,
np.mean(rewards), parameter1, parameter2, parameter3, parameter4,
parameter5
]]
path = os.path.join(results_path, 'results.csv')
if not os.path.isfile(path):
df = pd.DataFrame(data=data, index=[model_name], columns=columns)
else:
df = pd.read_csv(path, index_col=0)
if model_name not in df.index:
df = df.append(
pd.DataFrame(data=data, index=[model_name], columns=columns))
else:
df.loc[model_name] = data[0]
df.to_csv(path, index=True)
torch.save(rewards,
os.path.join(results_path, f'{model_name}_episode_rewards.pt'))