Skip to content
Snippets Groups Projects
Select Git revision
  • mia_panic_branch
  • master default
  • plotting
  • gen
  • testing_evo
  • 3daysago
  • 4daysago
  • 5daysago
  • Max
  • 6daysago
  • 0.as
  • 0.asdf
  • 0.01
  • 0.1
  • 0
15 results
Blame Permalink
parse_ds_cn_arch.py 3.14 KiB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64 






from MetaAugment.child_networks import *
from MetaAugment.main import create_toy, train_child_network
import torch
import torchvision
import torchvision.datasets as datasets
import pickle

def parse_ds_cn_arch(ds, ds_name, IsLeNet):
    if ds == "MNIST":
        train_dataset = datasets.MNIST(root='./datasets/mnist/train', 
                        train=True, download=True, transform=None)
        test_dataset = datasets.MNIST(root='./datasets/mnist/test', train=False, 
                        download=True, transform=torchvision.transforms.ToTensor())
    elif ds == "KMNIST":
        train_dataset = datasets.KMNIST(root='./datasets/kmnist/train', 
                        train=True, download=True, transform=None)
        test_dataset = datasets.KMNIST(root='./datasets/kmnist/test', train=False, 
                        download=True, transform=torchvision.transforms.ToTensor())
    elif ds == "FashionMNIST":
        train_dataset = datasets.FashionMNIST(root='./datasets/fashionmnist/train', 
                        train=True, download=True, transform=None)
        test_dataset = datasets.FashionMNIST(root='./datasets/fashionmnist/test', train=False, 
                        download=True, transform=torchvision.transforms.ToTensor())
    elif ds == "CIFAR10":
        train_dataset = datasets.CIFAR10(root='./datasets/cifar10/train', 
                        train=True, download=True, transform=None)
        test_dataset = datasets.CIFAR10(root='./datasets/cifar10/test', train=False, 
                        download=True, transform=torchvision.transforms.ToTensor())
    elif ds == "CIFAR100":
        train_dataset = datasets.CIFAR100(root='./datasets/cifar100/train', 
                        train=True, download=True, transform=None)
        test_dataset = datasets.CIFAR100(root='./datasets/cifar100/test', train=False, 
                        download=True, transform=torchvision.transforms.ToTensor())
    elif ds == 'Other':
        dataset = datasets.ImageFolder('./datasets/upload_dataset/'+ ds_name, transform=None)
        len_train = int(0.8*len(dataset))
        train_dataset, test_dataset = torch.utils.data.random_split(dataset, [len_train, len(dataset)-len_train])

        # check sizes of images
    img_height = len(train_dataset[0][0][0])
    img_width = len(train_dataset[0][0][0][0])
    img_channels = len(train_dataset[0][0])


        # check output labels
    if ds == 'Other':
        num_labels = len(dataset.class_to_idx)
    elif ds == "CIFAR10" or ds == "CIFAR100":
        num_labels = (max(train_dataset.targets) - min(train_dataset.targets) + 1)
    else:
        num_labels = (max(train_dataset.targets) - min(train_dataset.targets) + 1).item()


        
    if IsLeNet == "LeNet":
        child_architecture = LeNet(img_height, img_width, num_labels, img_channels)
    elif IsLeNet == "EasyNet":
        child_architecture = EasyNet(img_height, img_width, num_labels, img_channels)
    elif IsLeNet == 'SimpleNet':
        child_architecture = SimpleNet(img_height, img_width, num_labels, img_channels)
    else:
        child_architecture = pickle.load(open(f'datasets/childnetwork', "rb"))

    return train_dataset, test_dataset, child_architecture