import torch

from tqdm import tqdm

from torchvision import datasets, transforms, models

from torch.utils.data import DataLoader

torch.cuda.empty_cache()

model_names = [

        "cifar10_vgg11_bn",

        "cifar10_vgg13_bn",

        "cifar10_vgg16_bn",

        "cifar10_vgg19_bn",

        # "cifar10_resnet56",

]

batch_size = 1

test_transform = transforms.Compose([

                    transforms.ToTensor(),

])

def load_models():

    models = []

    for model_name in model_names:

        model = torch.hub.load("chenyaofo/pytorch-cifar-models", model_name, pretrained=True)

        models.append(model)

    return models

testset = datasets.CIFAR10(root='./data', train=False,

                                       download=True, transform=test_transform)

testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,

                                         shuffle=False, num_workers=1)

models = load_models()

model = torch.hub.load("chenyaofo/pytorch-cifar-models", "cifar10_resnet20", pretrained=True)

import torch.nn as nn
import torch

class MyEnsemble(nn.Module):

  def __init__(self, modelA, modelB, modelC, modelD):
    super(MyEnsemble,self).__init__()
    self.modelA = modelA
    self.modelB = modelB
    self.modelC = modelC
    self.modelD = modelD
    # self.modelE = modelE

  def forward(self, x):
    out1 = self.modelA(x)
    out2 = self.modelB(x)
    out3 = self.modelC(x)
    out4 = self.modelD(x)
    # out5 = self.modelE(x)

    out1 = torch.softmax(out1, dim=1)
    out2 = torch.softmax(out2, dim=1)
    out3 = torch.softmax(out3, dim=1)
    out4 = torch.softmax(out4, dim=1)
    # out5 = torch.softmax(out5, dim=1)

    out = out1 + out2 + out3 + out4

    # out = out / 2

    return out


device = 'cuda' if torch.cuda.is_available() else 'cpu'

model = MyEnsemble(models[0], models[1], models[2], models[3])

model.to(device)

total = 0
correct = 0
with torch.no_grad():
    for images, labels in tqdm(testloader):
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predictions = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predictions == labels).sum().item()

print('Accuracy of the network on the 10000 test images: %d %%' % (
    100 * correct / total))