def get_dataloader(self, sample_width, sample_height):
'''
Returns a train and a validate dataloader
'''
IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm',
'.tif', '.tiff', '.webp')
data_root = self.root_train_test_data
transformation = FileUtils.get_image_transforms(sample_width,
sample_height,
to_grayscale=False)
train_dataset = ImageFolder(
os.path.join(data_root, 'train'),
transformation,
is_valid_file=lambda file: Path(file).suffix in IMG_EXTENSIONS)
val_dataset = ImageFolder(
os.path.join(data_root, 'validation'),
transformation,
is_valid_file=lambda file: Path(file).suffix in IMG_EXTENSIONS)
train_loader = DataLoader(train_dataset,
batch_size=self.batch_size,
shuffle=True,
drop_last=True)
val_loader = DataLoader(val_dataset,
batch_size=self.batch_size,
shuffle=True,
drop_last=True)
return train_loader, val_loader
def __init__(self, path, batch_size, shuffle, num_workers, map_index_to_class):
self._path = path
self._image_folder = ImageFolder(self._path, transform=self.resize_and_to_tensor)
index_to_class_dictionary = self.swap_keys_values(self._image_folder.class_to_idx) if map_index_to_class else {}
self.index_to_class_dictionary = index_to_class_dictionary
self.filenames = self.parse_image_filenames(self._image_folder.imgs)
self.data_loader = DataLoader(self._image_folder, batch_size, shuffle=shuffle, num_workers=num_workers)
def get_finetune_dataloader(args): train_transforms = transforms.Compose([ transforms.Resize(256), transforms.RandomCrop(224), transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) val_transforms = transforms.Compose([ transforms.Resize(256), transforms.CenterCrop(224), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ]) train_dataset = ImageFolder(root='/data/Caltech101/train/', transform=train_transforms) val_dataset = ImageFolder(root='/data/Caltech101/val/', transform=val_transforms) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, pin_memory=False, num_workers=args.n_workers) val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, pin_memory=False, num_workers=args.n_workers) return train_loader, val_loader
def get_dataloaders(
train_dir,
var_dir,
train_transform=None,
val_transform=None,
split=(0.5, 0.5),
batch_size=32,
*args, **kwargs):
"""
This label_type returns the train, val and test dataloaders.
"""
# create the datasets
train_ds = ImageFolder(root=train_dir, transform=train_transform)
val_ds = ImageFolder(root=var_dir, transform=val_transform)
# now we want to split the val_ds in validation and test
lengths = np.array(split) * len(val_ds)
lengths = lengths.astype(int)
left = len(val_ds) - lengths.sum()
# we need to add the different due to float approx to int
lengths[-1] += left
val_ds, test_ds = random_split(val_ds, lengths.tolist())
logging.info(f'Train samples={len(train_ds)}, Validation samples={len(val_ds)}, Test samples={len(test_ds)}')
train_dl = DataLoader(train_ds, batch_size=batch_size, shuffle=True, *args, **kwargs)
val_dl = DataLoader(val_ds, batch_size=batch_size, shuffle=False, *args, **kwargs)
test_dl = DataLoader(test_ds, batch_size=batch_size, shuffle=False, *args, **kwargs)
return train_dl, val_dl, test_dl
def __init__(self, train=True, rootDir=None, transform=None, test=False):
super(TinyImageNet, self).__init__()
datasetURL = "http://cs231n.stanford.edu/tiny-imagenet-200.zip"
if rootDir is None:
rootDir = os.path.abspath(os.path.join(os.getcwd(), "../data"))
if not os.path.exists(
os.path.join(rootDir, "../data/tiny-imagenet-200")):
print(f"Downloading TinyImageNet data to {rootDir}")
download_and_extract_archive(datasetURL, rootDir)
print("...done")
self.rootDir = os.path.abspath(
os.path.join(rootDir, "tiny-imagenet-200"))
self.train = train
self.test = test
self.transforms = transforms
trainDataset = ImageFolder(os.path.join(self.rootDir, "train"),
transform)
testDataset = ImageFolder(os.path.join(self.rootDir, "test"),
transform)
validDataset = TinyImagenetVal(self.rootDir, transform)
if not self.test:
if self.train:
self._dataset = trainDataset
else:
self._dataset = validDataset
self.targets = self._dataset.targets
else:
self._dataset = testDataset
self.targets = None
def get_loaders(
data_dir,
train_transforms=None,
val_transforms=None,
train_test_split=0.85,
train_val_split=0.15,
batch_size=16,
shuffle=True):
"""
This function returns the training, validation, & test loaders.
"""
np.random.seed(24)
train_ds = ImageFolder(root=data_dir, transform=train_transforms)
val_ds = ImageFolder(root=data_dir, transform=val_transforms)
test_ds = ImageFolder(root=data_dir, transform=val_transforms)
img_count = len(train_ds)
indices = list(range(img_count))
test_split = int(img_count * train_test_split)
if shuffle:
np.random.shuffle(indices)
train_idx, test_idx = indices[:test_split], indices[test_split:]
train_count = len(train_idx)
val_split = int(train_count * (1 - train_val_split))
train_idx, val_idx = train_idx[:val_split], train_idx[val_split:]
train_sample = SubsetRandomSampler(train_idx)
val_sample = SubsetRandomSampler(val_idx)
test_sample = SubsetRandomSampler(test_idx)
train_loader = DataLoader(train_ds, batch_size=batch_size, sampler=train_sample)
val_loader = DataLoader(val_ds, batch_size=batch_size, sampler=val_sample)
test_loader = DataLoader(test_ds, batch_size=batch_size, sampler=test_sample)
return train_loader, val_loader, test_loader
def __init__(self,
root,
loader=default_loader,
transform=None,
target_transform=None):
super(Faces, self).__init__()
self.folder = ImageFolder(root,
transform=transform,
target_transform=target_transform,
loader=loader)
def __init__(
self,
root: str,
train: bool = True,
download: bool = False,
transform: Callable = None,
target_transform: Callable = None,
loader: Callable = default_loader,
use_default_transforms: bool = False,
):
ImageFolder.__init__(self, f"{root}/train", transform,
target_transform, loader)
BaseDataset.__init__(self, root, True, transform, target_transform,
False, use_default_transforms)
class ImageDataset:
def __init__(self, path, batch_size, shuffle, num_workers, map_index_to_class):
self._path = path
self._image_folder = ImageFolder(self._path, transform=self.resize_and_to_tensor)
index_to_class_dictionary = self.swap_keys_values(self._image_folder.class_to_idx) if map_index_to_class else {}
self.index_to_class_dictionary = index_to_class_dictionary
self.filenames = self.parse_image_filenames(self._image_folder.imgs)
self.data_loader = DataLoader(self._image_folder, batch_size, shuffle=shuffle, num_workers=num_workers)
# TODO add data augmentation to the transform
@staticmethod
def resize_and_to_tensor(pil_image):
return Compose([
Resize((224, 224)),
ToTensor()
])(pil_image)
@staticmethod
def parse_image_filenames(full_paths):
return list(map(lambda path: path[0].split('/')[-1], full_paths))
@staticmethod
def swap_keys_values(dictionary):
swapped_dictionary = {}
for k in dictionary:
v = dictionary[k]
swapped_dictionary[v] = k
return swapped_dictionary
def __len__(self):
return self._image_folder.__len__()
def _get_dataloaders(
input_size,
batch_size,
val_batch_size,
data_dir
):
pre_transform = transforms.Compose([
transforms.Resize(input_size),
transforms.CenterCrop(input_size)
])
post_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([0.6383, 0.5834, 0.5287], [0.2610, 0.2711, 0.2873])
])
validation_transform = transforms.Compose([
transforms.Resize(input_size),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
transforms.Normalize([0.6383, 0.5834, 0.5287], [0.2610, 0.2711, 0.2873])
])
logging.info("Get training dataloader")
trans_dataset = TransDatasetFolder(os.path.join(data_dir, 'training'), pre_transform=pre_transform, post_transform=post_transform)
train_dataloader = torch.utils.data.DataLoader(trans_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
logging.info("Get validation dataloader")
validation_dataset = ImageFolder(os.path.join(data_dir, 'validation'), validation_transform)
validation_dataloader = torch.utils.data.DataLoader(validation_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
return train_dataloader, validation_dataloader
def __getitem__(self, index):
img, target = ImageFolder.__getitem__(self, index)
#img = self.__image_transformer(img)
origin = img
img = transforms.ToPILImage()(img)
s = float(img.size[0]) / self.slice
a = s / 2
tiles = [None] * self.slice**2
for n in range(self.slice**2):
i = n // self.slice
j = n % self.slice
c = [a * i * 2 + a, a * j * 2 + a]
c = np.array([c[1] - a, c[0] - a, c[1] + a + 1,
c[0] + a + 1]).astype(int)
tile = img.crop(c.tolist())
tile = self.__augment_tile(tile)
# Normalize the patches indipendently to avoid low level features shortcut
#m, s = tile.view(3,-1).mean(dim=1).numpy(), tile.view(3,-1).std(dim=1).numpy()
#s[s==0]=1
#norm = transforms.Normalize(mean=m.tolist(),std=s.tolist())
#tile = norm(tile)
tiles[n] = tile
order = np.random.permutation(self.slice**2)
data = [tiles[order[t]] for t in range(self.slice**2)]
data = torch.stack(data, 0)
tiles = torch.stack(tiles)
return origin, data, order, tiles
def prepare_imagenet(args):
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
##### train transform #####
transform_list = [
# transforms.RandomCrop(224, padding=4),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]
train_transforms = transforms.Compose(transform_list)
##### test transform #####
transform_list = [
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]
test_transforms = transforms.Compose(transform_list)
root = 'data/ILSVRC2012'
if os.path.exists(root + '/cache.pth'):
print('Loading from cache')
train_dataset, test_dataset = torch.load(root + '/cache.pth')
else:
train_dataset = ImageFolder(root + '/train',
transform=train_transforms)
test_dataset = ImageFolder(root + '/val', transform=test_transforms)
print('Saving to cache')
torch.save((train_dataset, test_dataset), root + '/cache.pth')
train_loader = DataLoader(train_dataset,
batch_size=args['batch_size'],
shuffle=True,
num_workers=32,
drop_last=True)
test_loader = DataLoader(test_dataset,
batch_size=args['batch_size'] * 2,
shuffle=False,
num_workers=32)
return train_loader, test_loader
def get_labels(model, data_dir):
train_ds = ImageFolder(root=data_dir)
model.idx_to_class = train_ds.class_to_idx
model.idx_to_class = {
idx: label
for label, idx in model.idx_to_class.items()
}
return model.idx_to_class
class Faces(Dataset):
def __init__(self,
root,
loader=default_loader,
transform=None,
target_transform=None):
super(Faces, self).__init__()
self.folder = ImageFolder(root,
transform=transform,
target_transform=target_transform,
loader=loader)
def __getitem__(self, index):
return self.folder.__getitem__(index)
def __len__(self):
return self.folder.__len__()
def check(self, folder_path):
try:
dataset_pending = ImageFolder(root=folder_path)
for class_dir in os.listdir(folder_path):
if os.path.isfile(os.join(folder_path, class_dir)):
raise DATA_PATTERN_MISMATCH
except:
raise DATA_PATTERN_MISMATCH
return True
def create_dataset(path, data_augmentation):
"""
Convenience function for this project
:param str path: str path to root directory containing folders with samples for each class
:param data_augmentation: torch transforms object
:return: torch dataset
"""
dataset = ImageFolder(root=path, transform=data_augmentation)
return dataset
def get_recognition_dataset(item_source, resolution, increase_by_factor, item_usage):
if item_usage:
dataset = UsageBasedDataset(root=item_source, usage=item_usage, transform=get_recognition_transform(resolution))
else:
dataset = ImageFolder(root=item_source, transform=get_recognition_transform(resolution))
if increase_by_factor:
dataset = ConcatDataset([dataset for _ in range(increase_by_factor)])
return dataset
def data_preprocessing(self, chunk_dirs: Iterable, batch_size,
transform: callable) -> Iterable:
return DataLoader(dataset=ConcatDataset([
ImageFolder(root=chunk_dir, transform=transform)
for chunk_dir in chunk_dirs
]),
batch_size=batch_size,
shuffle=True,
pin_memory=False,
num_workers=10)
def predict_to_ensemble(model_name,
model_class,
model_state_pth,
image_size,
normalize,
nb_classes=15,
batch_size=15,
with_crops=True):
print(f'[+] predict {model_name}')
model = get_model(model_class, nb_classes, model_state_pth=model_state_pth)
model.eval()
tta_preprocess = [
preprocess(normalize, image_size),
preprocess_hflip(normalize, image_size)
]
if with_crops:
tta_preprocess += make_transforms(
[transforms.Resize((image_size + 20, image_size + 20))],
[transforms.ToTensor(), normalize], five_crops(image_size))
print(f'[+] tta size: {len(tta_preprocess)}')
data_loaders = []
for transform in tta_preprocess:
data_loader = get_data_loader('./test/',
data_transform=transform,
batch_size=batch_size)
data_loaders.append(data_loader)
lx, px = utils.predict_tta(model, data_loaders)
test_predict = {
'lx': lx.cpu(),
'px': px.cpu(),
}
torch.save(test_predict, f'{model_name}_test_prediction.pth')
data_loaders = []
for transform in tta_preprocess:
valid_dataset = ImageFolder('./data/train/', transform=transform)
data_loader = get_data_loader('./data/train/',
batch_size=batch_size,
dataset=valid_dataset)
data_loaders.append(data_loader)
lx, px = utils.predict_tta(model, data_loaders)
val_predict = {
'lx': lx.cpu(),
'px': px.cpu(),
}
torch.save(val_predict, f'{model_name}_val_prediction.pth')
return {'test': test_predict, 'val': val_predict}
def data_preprocessing(self, chunk_dirs: Iterable, batch_size,
transform: callable) -> Iterable:
data_loader = DataLoader(dataset=ConcatDataset([
ImageFolder(root=chunk_dir, transform=transform)
for chunk_dir in chunk_dirs
]),
batch_size=batch_size,
shuffle=True,
pin_memory=False,
num_workers=10,
collate_fn=fast_collate)
return DataPrefetcher(
data_loader,
normalize_mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
normalize_std=[0.229 * 255, 0.224 * 255, 0.225 * 255])
def get_dataloaders(train_dir,
val_dir,
test_dir,
train_transform=None,
val_transform=None,
batch_size=32,
*args,
**kwargs):
"""
This function returns the train, val and test dataloaders.
"""
# create the datasets
train_ds = ImageFolder(root=train_dir, transform=train_transform)
val_ds = ImageFolder(root=val_dir, transform=val_transform)
test_ds = ImageFolder(root=test_dir, transform=val_transform)
logging.info(
f'Train samples={len(train_ds)}, Validation samples={len(val_ds)}, Test samples={len(test_ds)}'
)
train_dl = DataLoader(train_ds,
batch_size=batch_size,
shuffle=True,
*args,
**kwargs)
val_dl = DataLoader(val_ds,
batch_size=batch_size,
shuffle=True,
*args,
**kwargs)
test_dl = DataLoader(test_ds,
batch_size=batch_size,
shuffle=True,
*args,
**kwargs)
return train_dl, val_dl, test_dl
def _build_dataset(self, data_rng: spec.RandomState, split: str,
data_dir: str, batch_size: int):
is_train = (split == "train")
normalize = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
mean=[i / 255 for i in self.train_mean],
std=[i / 255 for i in self.train_stddev])
])
transform_config = {
"train":
transforms.Compose([
transforms.RandomResizedCrop(
self.center_crop_size,
scale=self.scale_ratio_range,
ratio=self.aspect_ratio_range),
transforms.RandomHorizontalFlip(), normalize
]),
"test":
transforms.Compose([
transforms.Resize(self.resize_size),
transforms.CenterCrop(self.center_crop_size), normalize
])
}
folder = {'train': 'train', 'test': 'val'}
dataset = ImageFolder(
os.path.join(data_dir, folder[split]),
transform=transform_config[split])
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
shuffle=is_train,
num_workers=5,
pin_memory=True,
drop_last=is_train)
if is_train:
dataloader = cycle(dataloader)
return dataloader
def make_datasets(dataset_path: Path, dev_ratio=0.1, test_ratio=0.1, rng=None):
if rng is None:
rng = np.random.default_rng()
# find images
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
transforms.Resize((224, 224))
])
dataset = ImageFolder(dataset_path)
imgs = dataset.imgs
labels = dataset.targets
label_indices = defaultdict(list)
for i, label in enumerate(labels):
label_indices[label].append(i)
train_indices = []
dev_indices = []
test_indices = []
for label, indices in label_indices.items():
n = len(indices)
n_dev = int(math.ceil(n * dev_ratio))
n_test = int(math.ceil(n * test_ratio))
n_train = n - (n_dev + n_test)
rng.shuffle(indices)
train_indices.extend(indices[:n_train])
dev_indices.extend(indices[n_train:n_train + n_dev])
test_indices.extend(indices[n_train + n_dev:])
train_dataset = ImageDatasetFromLists([imgs[i] for i in train_indices],
[labels[i] for i in train_indices],
root=dataset_path,
transform=transform)
dev_dataset = ImageDatasetFromLists([imgs[i] for i in dev_indices],
[labels[i] for i in dev_indices],
root=dataset_path,
transform=transform)
test_dataset = ImageDatasetFromLists([imgs[i] for i in test_indices],
[labels[i] for i in test_indices],
root=dataset_path,
transform=transform)
return train_dataset, dev_dataset, test_dataset
def get_data_loader(path,
data_transform=None,
batch_size=1,
num_workers=1,
dataset=None):
if dataset is None:
if data_transform is None:
raise
test_dataset = ImageFolder('./test/', data_transform)
else:
test_dataset = dataset
test_loader = DataLoader(
test_dataset,
batch_size=batch_size,
sampler=None,
num_workers=num_workers,
pin_memory=True,
)
return test_loader
def write_img_grid(
cls,
writer,
img_root_dir,
num_imgs=4,
class_sample_file_pairs=None,
img_height=200, # px
img_width=400, # px
to_grayscale=True,
unittesting=False):
'''
Create and log a Tensorboard 'grid' of
example train images. The img_root_dir must
be the 'data root': the dir holding one subdir
per class.
:param writer: a Tensorboard Pytorch SummaryWriter
:type writer: SummaryWriter
:param img_root_dir: directory
that contains sub-directories with samples. The
sub-directory names are taken to be class names.
:type img_root_dir: str
:param num_imgs: total number of images to
include in the grid. If None: all images
:type num_imgs: {None | int}
:param class_sample_file_pairs: <class>/<img_file_name> for
individual images if random choice is not wanted.
:type class_sample_file_pairs: {None | str | [str]}
:param img_height: height of all images
:type img_height: int (pixels)
:param img_width: width of all images
:type img_width: int (pixels)
:param to_grayscale: whether or not to convert
images to grayscale upon import
:type to_grayscale: bool
:param unittesting: controls whether grid is
actually created, or the img tensor that
would be contained in the grid is returned
for testing dimensions.
:type unittesting: bool
'''
if img_root_dir is None:
raise ValueError("Must provide path to image root dir")
# Prepare to resize all images to a given dimension,
# convert to grayscale if requested, and turn into
# a tensor:
the_transforms = [transforms.Resize((img_height, img_width))]
if to_grayscale:
the_transforms.append(transforms.Grayscale())
the_transforms.append(transforms.ToTensor())
transform_img = transforms.Compose(the_transforms)
# Get an ImageFolder instance, from which
# we will easily find classes and samples
img_folder = ImageFolder(img_root_dir,
transform=transform_img,
loader=default_loader)
# Get list of full paths to samples:
sample_idxs = cls._get_sample_indices(
img_folder,
num_imgs=num_imgs,
class_sample_file_pairs=class_sample_file_pairs)
# Get list of img tensor/class_idx pairs:
img_tns_list = [img_folder[idx] for idx in sample_idxs]
# Print <class>/file_name onto
# each spectrogram:
marked_img_tns_list = []
for i, (img_tns, class_idx) in enumerate(img_tns_list):
class_name = img_folder.classes[class_idx]
# img_folder.samples is [ (full_path, class_idx), (..., ...) ]:
img_file_basename = os.path.basename(img_folder.samples[i][0])
marked_img_tns_list.append(
cls.print_onto_image(img_tns,
f"{class_name}/{img_file_basename}"))
# Turn list of img tensors into
# a single tensor with first dim
# being len of list:
marked_img_tns = torch.cat(marked_img_tns_list)
# A 10px frame around each img:
grid = make_grid(marked_img_tns, padding=10)
if unittesting:
return grid
writer.add_image('Train Input Examples', grid)
return grid
# The network is developed by Visual Geometry Group of Univ. of Oxford, so named 'VGGNet'
import torch.nn as nn
from torch.utils.data.dataloader import DataLoader
import torchvision.transforms.transforms as transforms
from torchvision.datasets.folder import ImageFolder
from Nets1 import VGGNet
# 1. Load the data
transform = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
testset = ImageFolder('C:/Users/Cooper/PycharmProjects/Imagenet', transform=transform)
testloader = DataLoader(testset, batch_size=10)
# 2. Define a network
net = VGGNet('C')
# 4. Test the network
for img, label in testloader:
output = net(img)
_, predicted = output.max(1)
print(predicted)
def str2dataset(name, device="cuda", train=False):
if name == "MNIST" or name == "mnist":
mnist = torchvision.datasets.MNIST(root="./data",
train=train,
download=True,
transform=transforms.ToTensor())
return (mnist, lambda x: x, lambda x: x)
elif name == "CIFAR" or name == "cifar":
if train:
transform = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
])
else:
transform = transforms.ToTensor()
cifar = torchvision.datasets.CIFAR10(root="./data",
train=train,
download=True,
transform=transform)
mu = torch.tensor([0.4914, 0.4822, 0.4465],
dtype=torch.float,
device=device).unsqueeze(-1).unsqueeze(-1)
std = torch.tensor([0.2023, 0.1994, 0.2010],
dtype=torch.float,
device=device).unsqueeze(-1).unsqueeze(-1)
normalize = lambda x: (x - mu) / std
unnormalize = lambda x: x * std + mu
return (cifar, normalize, unnormalize)
elif name == "ImageNet" or name == "imagenet":
imagenet = ImageNetLoader(root="./data/ImageNet",
split="val",
transform=transforms.Compose([
transforms.CenterCrop(size=224),
transforms.ToTensor()
]))
mu = torch.tensor([0.485, 0.456, 0.406],
dtype=torch.float,
device=device).unsqueeze(-1).unsqueeze(-1)
std = torch.tensor([0.229, 0.224, 0.225],
dtype=torch.float,
device=device).unsqueeze(-1).unsqueeze(-1)
normalize = lambda x: (x - mu) / std
unnormalize = lambda x: x * std + mu
return (imagenet, normalize, unnormalize)
elif name == "Toy" or name == "toy":
toyset = torch.utils.data.TensorDataset(
torch.tensor([[0, 1], [1, 0]], dtype=torch.float).view(2, 1, 1, 2),
torch.tensor([1, 0], dtype=torch.long))
return (toyset, lambda x: x, lambda x: x)
elif name == "Yale" or name == "yale":
yale_faces = ImageFolder("./yale_data",
transform=transforms.Compose([
transforms.CenterCrop(size=224),
transforms.ToTensor()
]))
return (yale_faces, lambda x: x, lambda x: x)
else:
raise Exception('data set not supported')
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.utils.data.dataloader import default_collate
# 计算机视觉的数据读取类
from torchvision.datasets.folder import ImageFolder, default_loader
import pandas as pd
class MyDataset(Dataset):
def __init__(self, csv_file, txt_file, root_dir, other_file):
self.csv_data = pd.read_csv(csv_file)
with open(txt_file, 'r') as f:
data_list = f.readlines()
self.txt_data = data_list
self.root_dir = root_dir
def __len__(self):
return len(self.csv_data)
def __getitem__(self, idx):
data = (self.csv_data[idx], self.txt_data[idx])
return data
dataiter = DataLoader(MyDataset,
batch_size=32,
shuffle=True,
collate_fn=default_collate)
dset = ImageFolder(root='root_path', transform=None, loader=default_loader)
def train():
train_dataset = ImageFolder('./data/train/', transform=preprocess_with_augmentation(normalize_torch, IMAGE_SIZE))
valid_dataset = ImageFolder('./data/train/', transform=preprocess(normalize_torch, IMAGE_SIZE))
training_data_loader, valid_data_loader = (split_train_val_loader(train_dataset, valid_dataset,
len(train_dataset), valid_size=VALID_SIZE, batch_size=BATCH_SIZE,
train_enlarge_factor=TRAIN_ENLARGE_FACTOR,
pin_memory=True, num_workers=1, random_seed=RANDOM_SEED
))
model = get_model(MODEL, NB_CLASSES)
criterion = nn.CrossEntropyLoss().cuda()
nb_learnable_params = sum(p.numel() for p in model.fresh_params())
print(f'[+] nb learnable params {nb_learnable_params}')
lx, px = utils.predict(model, valid_data_loader, prob=False)
min_loss = criterion(Variable(px), Variable(lx)).item()
_, preds = torch.max(px.data, dim=1)
accuracy = torch.mean((preds != lx).float())
print(f' original loss: {min_loss}, accuracy: {accuracy}')
lr = 0.001
patience = 0
earlystop = 0
optimizer = torch.optim.Adam(model.fresh_params(), lr=lr)
torch.save(model.state_dict(), MODEL_FILE_NAME)
for epoch in range(EPOCH):
if epoch == 1:
lr = 0.0005
print(f'[+] set lr={lr}')
if patience == PATIENCE_LIMIT:
patience = 0
model.load_state_dict(torch.load(MODEL_FILE_NAME))
lr = lr / 10
print(f'[+] set lr={lr}')
if earlystop > EARLY_STOP:
model.load_state_dict(torch.load(MODEL_FILE_NAME))
print('EARLY STOPPED')
continue
if epoch > 0:
optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=0.0001)
running_loss = RunningMean()
running_score = RunningMean()
model.train()
pbar = tqdm(training_data_loader, total=len(training_data_loader))
for inputs, labels in pbar:
batch_size = inputs.size(0)
inputs = Variable(inputs)
labels = Variable(labels)
if use_gpu:
inputs = inputs.cuda()
labels = labels.cuda()
optimizer.zero_grad()
outputs = model(inputs)
_, preds = torch.max(outputs.data, dim=1)
loss = criterion(outputs, labels)
running_loss.update(loss.item(), 1)
running_score.update(torch.sum(preds == labels.data).float(), batch_size)
loss.backward()
optimizer.step()
pbar.set_description(f'{epoch}: {running_loss.value:.5f} {running_score.value:.3f}')
model.eval()
lx, px = utils.predict(model, valid_data_loader)
log_loss = criterion(Variable(px), Variable(lx))
log_loss = log_loss.item()
_, preds = torch.max(px, dim=1)
accuracy = torch.mean((preds == lx).float())
print(f'[+] val loss: {log_loss:.5f} acc: {accuracy:.3f}')
if (log_loss < min_loss):
torch.save(model.state_dict(), MODEL_FILE_NAME)
print(f'[+] val loss improved from {min_loss:.5f} to {log_loss:.5f}, accuracy={accuracy}. Saved!')
min_loss = log_loss
patience = 0
else:
patience += 1
earlystop += 1
def get_imagenet_train_folder(path):
return ImageFolder(path, transform=train_transforms, loader=default_loader)
