def getNetwork(args,inputs,outputs):
if (args.net_type == 'lenet'):
net = BBBLeNet(outputs,inputs) # inputs is number of input channels
file_name = 'lenet'
elif (args.net_type == 'alexnet'):
net = BBBAlexNet(outputs,inputs)
file_name = 'alexnet-'
elif (args.net_type == '3conv3fc'):
net = BBB3Conv3FC(outputs,inputs)
file_name = '3Conv3FC-'
else:
print('Error : Network should be either [LeNet / AlexNet /SqueezeNet/ 3Conv3FC')
sys.exit(0)
return net, file_name
ax.annotate(str(val),
xy=(y, x),
horizontalalignment='center',
verticalalignment='center',
size=8,
color='white' if img[x][y] < thresh else 'black')
# In[ ]:
# Architecture
if (net_type == 'lenet'):
net = BBBLeNet(outputs, inputs)
elif (net_type == 'alexnet'):
net = BBBAlexNet(outputs, inputs)
elif (net_type == '3conv3fc'):
net = BBB3Conv3FC(outputs, inputs)
else:
print('Error : Network should be either [LeNet / AlexNet / 3Conv3FC')
# In[ ]:
if use_cuda:
net.cuda()
# In[ ]:
vi = GaussianVariationalInference(torch.nn.CrossEntropyLoss())
optimizer = optim.Adam(net.parameters(), lr=lr, weight_decay=weight_decay)
# In[ ]:
