import os import sys import math import pytest import unittest import numpy as np import torch import torch.nn as nn

import utils from models.BayesianModels.Bayesian3Conv3FC import BBB3Conv3FC from models.BayesianModels.BayesianAlexNet import BBBAlexNet from models.BayesianModels.BayesianLeNet import BBBLeNet from models.NonBayesianModels.AlexNet import AlexNet from models.NonBayesianModels.LeNet import LeNet from models.NonBayesianModels.ThreeConvThreeFC import ThreeConvThreeFC from layers.BBBConv import BBBConv2d from layers.BBBLinear import BBBLinear from layers.misc import FlattenLayer

cuda_available = torch.cuda.is_available() bayesian_models = [BBBLeNet, BBBAlexNet, BBB3Conv3FC] non_bayesian_models = [LeNet, AlexNet, ThreeConvThreeFC]

class TestModelForwardpass:

@pytest.mark.parametrize("model", bayesian_models)
def test_cpu_bayesian(self, model):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 3, 32, 32))
    net = model(10, 3)
    out = net(batch)
    assert out[0].shape[0]==batch_size

@pytest.mark.parametrize("model", non_bayesian_models)
def test_cpu_frequentist(self, model):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 3, 32, 32))
    net = model(10, 3)
    out = net(batch)
    assert out.shape[0]==batch_size

@pytest.mark.skipif(not cuda_available, reason="CUDA not available")
@pytest.mark.parametrize("model", bayesian_models)
def test_gpu_bayesian(self, model):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 3, 32, 32))
    net = model(10, 3)
    if cuda_available:
        net = net.cuda()
        batch = batch.cuda()
    out = net(batch)
    assert out[0].shape[0]==batch_size

@pytest.mark.skipif(not cuda_available, reason="CUDA not available")
@pytest.mark.parametrize("model", non_bayesian_models)
def test_gpu_frequentist(self, model):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 3, 32, 32))
    net = model(10, 3)
    if cuda_available:
        net = net.cuda()
        batch = batch.cuda()
    out = net(batch)
    assert out.shape[0]==batch_size

class TestBayesianLayers:

def test_flatten(self):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 64, 4, 4))

    layer = FlattenLayer(4 * 4 * 64)
    batch = layer(batch)

    assert batch.shape[0]==batch_size
    assert batch.shape[1]==(4 * 4 *64)

def test_conv(self):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 16, 24, 24))

    layer = BBBConv2d(16, 6, 4, alpha_shape=(1,1), padding=0, bias=False)
    batch = layer(batch)

    assert batch.shape[0]==batch_size
    assert batch.shape[1]==6

def test_linear(self):
    batch_size = np.random.randint(1, 256)
    batch = torch.randn((batch_size, 128))

    layer = BBBLinear(128, 64, alpha_shape=(1,1), bias=False)
    batch = layer(batch)

    assert batch.shape[0]==batch_size
    assert batch.shape[1]==64