MD-SNN/load_settings.py at main · Intelligent-Computing-Lab-Panda/MD-SNN · GitHub

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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import torch
import torchvision
from torch.utils.data.distributed import DistributedSampler

import os
import models.MS_ResNet as MS
import models.Quan_MS_ResNet as QMS
import numpy as np
import random


def load_paper_settings(args):
    t_window = args.time_window
    depth = args.model_depth

    if args.data == 'CIFAR10':
        num_class = 10
    elif args.data == 'CIFAR100':
        num_class = 100
    elif args.data == 'TinyImagenet':
        num_class = 200
    elif args.data == 'ncaltech101':
        num_class = 101
    else:
        raise ValueError("Unsupported dataset: {}".format(args.data))

    args.nc = num_class

    if args.paper_setting == 'student':
        print('Only quantization model')
        teacher = None
        if depth == 18:
            student = QMS.qresnet18(args, t_window, num_class)
        elif depth == 34:
            student = QMS.qresnet34(args, t_window, num_class)
        else:
            raise ValueError("Unsupported model_depth: {}".format(depth))

    elif args.paper_setting == 'original':
        print('Original model')
        teacher = None
        if depth == 18:
            student = MS.resnet18(args, t_window, num_class)
        elif depth == 34:
            student = MS.resnet34(args, t_window, num_class)
        else:
            raise ValueError("Unsupported model_depth: {}".format(depth))

    elif args.paper_setting == 'qamd':
        print('KD model')
        if depth == 18:
            teacher = MS.resnet18(args, t_window, num_class)
            student = QMS.qresnet18(args, t_window, num_class)
        elif depth == 34:
            teacher = MS.resnet34(args, t_window, num_class)
            student = QMS.qresnet34(args, t_window, num_class)
        else:
            raise ValueError("Unsupported model_depth: {}".format(depth))

    else:
        raise ValueError("Unsupported paper_setting: {}".format(args.paper_setting))

    return teacher, student, args


def generate_dataloader(args, num_gpus, data, name, transform):
    if data is None:
        return None

    if transform is None:
        dataset = torchvision.datasets.ImageFolder(data, transform=torchvision.transforms.ToTensor())
    else:
        dataset = torchvision.datasets.ImageFolder(data, transform=transform)

    use_cuda = torch.cuda.is_available()
    kwargs = {"pin_memory": True, "num_workers": 4} if use_cuda else {}

    if name == "train":
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=args.batch_size // num_gpus,
            shuffle=False, sampler=DistributedSampler(dataset), **kwargs)
    elif name == "val":
        dataloader = torch.utils.data.DataLoader(
            dataset, batch_size=args.batch_size, shuffle=False, **kwargs)

    return dataloader


def seed_all(seed=1029):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.benchmark = False
    torch.backends.cudnn.deterministic = True