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
Пример #2
0
        for y in range(height):
            val = round(img[x][y], 2) if img[x][y] != 0 else 0
            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)