/
utils.py
94 lines (79 loc) · 3.41 KB
/
utils.py
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
import numpy
import torch
from torch.optim.lr_scheduler import _LRScheduler
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import torch.utils.data.sampler as sampler
from dataset import *
def get_training_dataloader(mean, std, batch_size=16, num_workers=2, sampler=sampler):
""" return training dataloader
Args:
mean: mean of cifar10 training dataset
std: std of cifar10 training dataset
path: path to cifar10 training python dataset
batch_size: dataloader batchsize
num_workers: dataloader num_works
shuffle: whether to shuffle
Returns: train_data_loader:torch dataloader object
"""
transform_train = transforms.Compose([
#transforms.ToPILImage(),
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
cifar10_training = CIFAR10Train(path='./data', transform=transform_train)
cifar10_training_loader = DataLoader(
cifar10_training, sampler=sampler, num_workers=num_workers, batch_size=batch_size, drop_last=True)
return cifar10_training_loader
def get_test_dataloader(mean, std, batch_size=16, num_workers=2, shuffle=True):
""" return training dataloader
Args:
mean: mean of cifar100 test dataset
std: std of cifar100 test dataset
path: path to cifar100 test python dataset
batch_size: dataloader batchsize
num_workers: dataloader num_works
shuffle: whether to shuffle
Returns: cifar100_test_loader:torch dataloader object
"""
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean, std)
])
# cifar100_test = CIFAR100Test(path='./data', transform=transform_test)
cifar10_test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform_test)
cifar10_test_loader = DataLoader(
cifar10_test, shuffle=shuffle, num_workers=num_workers, batch_size=batch_size)
return cifar10_test_loader
def compute_mean_std(cifar10_dataset):
"""compute the mean and std of cifar100 dataset
Args:
cifar100_training_dataset or cifar100_test_dataset
witch derived from class torch.utils.data
Returns:
a tuple contains mean, std value of entire dataset
"""
data_r = numpy.dstack([cifar10_dataset[i][1][:, :, 0] for i in range(len(cifar10_dataset))])
data_g = numpy.dstack([cifar10_dataset[i][1][:, :, 1] for i in range(len(cifar10_dataset))])
data_b = numpy.dstack([cifar10_dataset[i][1][:, :, 2] for i in range(len(cifar10_dataset))])
mean = numpy.mean(data_r), numpy.mean(data_g), numpy.mean(data_b)
std = numpy.std(data_r), numpy.std(data_g), numpy.std(data_b)
return mean, std
class WarmUpLR(_LRScheduler):
"""warmup_training learning rate scheduler
Args:
optimizer: optimzier(e.g. SGD)
total_iters: totoal_iters of warmup phase
"""
def __init__(self, optimizer, total_iters, last_epoch=-1):
self.total_iters = total_iters
super().__init__(optimizer, last_epoch)
def get_lr(self):
"""we will use the first m batches, and set the learning
rate to base_lr * m / total_iters
"""
return [base_lr * self.last_epoch / (self.total_iters + 1e-8) for base_lr in self.base_lrs]