def load_data(dataset_name="car_dataset",
input_size=224,
batch_size=32,
data_dir="./data"):
# Data augmentation and normalization for training
# Just normalization for validation
data_transforms = {
'train':
transforms.Compose([
transforms.RandomResizedCrop(input_size),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
'val':
transforms.Compose([
transforms.Resize(input_size),
transforms.CenterCrop(input_size),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
print("Initializing Datasets and Dataloaders...")
# check cuda availability
kwargs = {
'num_workers': 1,
'pin_memory': True
} if torch.cuda.is_available() else {}
# Create training and validation datasets
if dataset_name == "car_data":
if os.path.exists(os.path.join(
data_dir + '/car_data', 'train')) and os.path.exists(
os.path.join(data_dir + '/car_data', 'val')):
image_datasets = {
x:
datasets.ImageFolder(os.path.join(data_dir + '/car_data', x),
data_transforms[x])
for x in ['train', 'val']
}
else:
print(
"Please manually download the car dataset before training model\n"
)
print(
"Download Link: https://ai.stanford.edu/~jkrause/cars/car_dataset.html"
)
exit()
elif dataset_name == "car_data_modified":
if os.path.exists(
os.path.join(data_dir + '/car_data_modified',
'train')) and os.path.exists(
os.path.join(data_dir + '/car_data_modified',
'val')):
image_datasets = {
x: datasets.ImageFolder(
os.path.join(data_dir + '/car_data_modified', x),
data_transforms[x])
for x in ['train', 'val']
}
else:
print(
"Please manually download the car dataset before training model\n"
)
print(
"Download Link: https://drive.google.com/file/d/11bS7Az-x4WkMUM066KgAhyiVWx-BqGwa/view?usp=sharing"
)
exit()
# Load the cifar100 dataset from TorchVision
elif dataset_name == "cifar100":
image_datasets = {
"train":
datasets.CIFAR100('./data',
train=True,
transform=transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
]),
download=True),
"val":
datasets.CIFAR100('./data',
train=False,
transform=transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
]))
}
# Load the mnist dataset from TorchVision
elif dataset_name == "mnist":
image_datasets = {
"train":
datasets.MNIST('./data',
train=True,
download=True,
transform=transforms.Compose([
transforms.Resize(224),
transforms.Grayscale(3),
transforms.ToTensor(),
])),
"val":
datasets.MNIST('./data',
train=False,
transform=transforms.Compose([
transforms.Resize(224),
transforms.Grayscale(3),
transforms.ToTensor(),
]))
}
else:
print("UNKNOWN Dataset! Please Chooses from the following datasets:")
print("\t[car_data, car_data_modified, cifar100, mnist]")
exit()
classes = image_datasets["train"].classes
class_names = {i: name for i, name in enumerate(classes)}
# print(class_names)
# print(image_datasets["val"].targets)
# Create training and validation dataloaders
dataloaders_dict = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True,
**kwargs)
for x in ['train', 'val']
}
inputs, labels = next(iter(dataloaders_dict['train']))
return dataloaders_dict, class_names, image_datasets, classes
transforms.Resize((256, 128), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
if opt.PCB:
data_transforms = transforms.Compose([
transforms.Resize((384, 192), interpolation=3),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
# dataset_list = ['gallery', 'query', 'train_all']
dataset_list = ['gallery', 'query']
image_datasets = {
x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms)
for x in dataset_list
}
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=opt.batchsize,
shuffle=False,
num_workers=16)
for x in dataset_list
}
class_names = image_datasets[dataset_list[1]].classes
use_gpu = torch.cuda.is_available()
######################################################################
# Extract feature
def main_worker(gpu, ngpus_per_node, args):
global best_acc1
args.gpu = None
if args.gpu is not None:
print("Use GPU: {} for training".format(args.gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
world_size=args.world_size, rank=args.rank)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(
(args.workers + ngpus_per_node - 1) / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
#model.cuda()
# DistributedDataParallel will divide and allocate batch_size to all
# available GPUs if device_ids are not set
model = torch.nn.parallel.DistributedDataParallel(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
#model.cuda()
else:
model = torch.nn.DataParallel(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss() #.cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(), args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
if args.gpu is None:
checkpoint = torch.load(args.resume)
else:
# Map model to be loaded to specified single gpu.
loc = 'cuda:{}'.format(args.gpu)
checkpoint = torch.load(args.resume, map_location=loc)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
if args.data != None:
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
pass
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
transform = transforms.Compose([
# you can add other transformations in this list
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])
if args.data == None:
train_dataset = datasets.FakeData(size=100000000, image_size=(
3, 224, 224), num_classes=200, transform=transform)
val_loader = torch.utils.data.DataLoader(datasets.FakeData(
size=1001, image_size=(3, 224, 224), num_classes=200, transform=transform))
pass
else:
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
val_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(valdir, transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size, shuffle=False,
num_workers=args.workers, pin_memory=True)
pass
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(
train_dataset, batch_size=args.batch_size, shuffle=(
train_sampler is None),
num_workers=args.workers, pin_memory=True, sampler=train_sampler)
if args.evaluate:
validate(val_loader, model, criterion, args)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
if args.so_one_shot:
sys.stdout.flush()
return
def main():
global args, best_prec1
args = parser.parse_args()
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
# model = models.vgg16(pretrained=True)
model = NetworkNew_test('checkpoint/model.pth')
print(model)
model = torch.nn.DataParallel(model.cuda(),
device_ids=range(torch.cuda.device_count()))
cudnn.benchmark = True
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda()
optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,
model.parameters()),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Data loading code from lmdb
if args.load_from_lmdb:
train_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-train.lmdb'), True),
batch_size=args.batch_size,
num_workers=16,
pin_memory=True)
print('train_loader_success!')
val_loader = torch.utils.data.DataLoader(lmdbDataset(
os.path.join(args.data_base, 'ILSVRC-val.lmdb'), False),
batch_size=args.batch_size,
num_workers=8,
pin_memory=True)
else:
traindir = os.path.join('/opt/luojh/Dataset/ImageNet/images', 'train')
valdir = os.path.join('/opt/luojh/Dataset/ImageNet/images', 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
# evaluate and train
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
if is_best:
folder_path = 'checkpoint/fine_again'
if not os.path.exists(folder_path):
os.makedirs(folder_path)
best_prec1 = max(prec1, best_prec1)
torch.save(model.state_dict(), folder_path + '/model.pth')
print('best accuracy is %.3f' % best_prec1)
def __init__(self, options, path):
"""Prepare the network, criterion, solver, and data.
Args:
options, dict: Hyperparameters.
"""
print('Prepare the network and data.')
self._options = options
self._path = path
# pdb.set_trace()
# Network.
self._net = torch.nn.DataParallel(BCNN()).cuda()
# Load the model from disk.
# self._net.load_state_dict(self.load_my_state_dict(torch.load(self._path['model'])))
# self._net.load_state_dict(torch.load(self._path['model']))
print(self._net)
# Criterion.
self._criterion = torch.nn.CrossEntropyLoss().cuda()
# Solver.
self._solver = torch.optim.SGD(
self._net.module.fc.parameters(),
lr=self._options['base_lr'],
momentum=0.9,
weight_decay=self._options['weight_decay'])
self._scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
self._solver,
mode='max',
factor=0.1,
patience=5,
verbose=True,
threshold=1e-4)
train_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=448), # Let smaller edge match
torchvision.transforms.RandomHorizontalFlip(),
torchvision.transforms.RandomCrop(size=448),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
test_transforms = torchvision.transforms.Compose([
torchvision.transforms.Resize(size=448),
torchvision.transforms.CenterCrop(size=448),
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))
])
data_dir1 = '/cvdata/xuelu/CUB_200_2011/train'
train_data = datasets.ImageFolder(root=data_dir1,
transform=train_transforms)
data_dir2 = '/cvdata/xuelu/CUB_200_2011/test'
test_data = datasets.ImageFolder(root=data_dir2,
transform=test_transforms)
self._train_loader = torch.utils.data.DataLoader(
train_data,
batch_size=self._options['batch_size'],
shuffle=True,
num_workers=4,
pin_memory=True)
self._test_loader = torch.utils.data.DataLoader(test_data,
batch_size=8,
shuffle=False,
num_workers=4,
pin_memory=True)
help='number of workers for dataloader')
parser.add_argument('-b',
type=int,
default=16,
help='batch size for dataloader')
parser.add_argument('-s',
type=bool,
default=True,
help='whether shuffle the dataset')
args_dict = vars(parser.parse_args())
logger.info(args_dict)
net_type = args_dict['net']
# device = torch.device("cpu")
device = torch.device("gpu")
list_author = next(os.walk(EXAMPLE_FOLDER))[1] # 53
num_classes = len(list_author)
net = build_network(archi=net_type, use_gpu=False, num_classes=num_classes)
logger.info(net)
# net.load_state_dict(torch.load(args_dict['weights'], map_location=torch.device('cpu')))
net.load_state_dict(
torch.load(args_dict['weights'], map_location=torch.device('gpu')))
net.eval()
example_image_dir = './example_dataset_structure'
dataset = datasets.ImageFolder(example_image_dir, transform=None)
idx_to_class = {v: k for k, v in dataset.class_to_idx.items()}
app.run(host='0.0.0.0', port=8889)
def main():
data_dir = "flowers/"
train_dir = data_dir + 'train/'
valid_dir = data_dir + 'valid/'
test_dir = data_dir + 'test/'
#Defining transforms for the training, validation, and testing sets
train_transforms = transforms.Compose([transforms.RandomRotation(30),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
test_transforms = transforms.Compose([transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
validation_transforms = transforms.Compose([transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])])
# Loading the datasets with ImageFolder
train_dataset = datasets.ImageFolder(train_dir, transform = train_transforms)
test_dataset = datasets.ImageFolder(test_dir, transform = test_transforms)
validation_dataset = datasets.ImageFolder(valid_dir, transform = validation_transforms)
# Using the image datasets and the trainforms, define the dataloaders
trainloader = torch.utils.data.DataLoader(train_dataset,batch_size = 64, shuffle = True)
testloader = torch.utils.data.DataLoader(test_dataset,batch_size = 64)
validationloader = torch.utils.data.DataLoader(validation_dataset,batch_size = 64)
"""Loading Data Completed"""
""" Now loading checkpoint"""
# TODO: Write a function that loads a checkpoint and rebuilds the model
model = checkpoint_loader("save_directory/checkpoint.pth")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr = 0.002)
print("\n ...Testing Initiated ... \n")
test_loss = 0
accuracy = 0
steps = 0
model.eval()
with torch.no_grad():
for inputs, labels in testloader:
steps = steps+1
print(steps, "/", len(testloader))
inputs, labels = inputs.to(device), labels.to(device)
logps = model(inputs) #or model.forward(inputs)
batch_loss = criterion(logps,labels)
test_loss += batch_loss.item()
#accuracy
ps = torch.exp(logps)
top_p, top_class = ps.topk (1, dim =1)
equals = top_class == labels.view(*top_class.shape)
accuracy += torch.mean(equals.type(torch.FloatTensor)).item()
print(f"current accuracy: {accuracy/len(testloader):.3f}")
print("...done testing...")
print(f"Test loss: {test_loss/len(testloader):.3f}.."
f"Test accuracy: {accuracy/len(testloader):.3f}"
"\n ...Testing Stopped... \n")
model.train()
print("I'm done with the training")
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
use_gpu = True and torch.cuda.is_available()
train_dir = domainData['amazon'] # 'amazon', 'dslr', 'webcam'
val_dir = domainData['webcam']
num_classes = NUM_CLASSES['office']
print("use gpu: ", use_gpu)
torch.manual_seed(7)
if use_gpu:
torch.cuda.manual_seed(7)
image_datasets = {
'train': datasets.ImageFolder(train_dir, data_transforms['train']),
'val': datasets.ImageFolder(val_dir, data_transforms['val'])
}
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=32,
shuffle=True,
num_workers=4)
for x in ['train', 'val']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
def train_model(model,
clscriterion,
transform_train = transforms.Compose([
transforms.Resize((224, 224)),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(10),
transforms.RandomAffine(0, shear=10, scale=(0.8, 1.2)),
transforms.ColorJitter(brightness=1, contrast=1, saturation=1),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
train_ds = datasets.ImageFolder('ants_and_bees/train',
transform=transform_train)
train_loader = torch.utils.data.DataLoader(dataset=train_ds,
batch_size=20,
shuffle=True)
validation_ds = datasets.ImageFolder('ants_and_bees/val', transform=transform)
validation_loader = torch.utils.data.DataLoader(dataset=validation_ds,
batch_size=20)
classes = ['ant', 'bee']
model = models.alexnet(False)
model.load_state_dict(torch.load('alexnet-owt-4df8aa71.pth'))
for param in model.features.parameters():
param.requires_grad = False
n_inputs = model.classifier[6].in_features
def __init__(self,
save_path=None,
train_batch_size=256,
test_batch_size=512,
valid_size=None,
n_worker=32,
resize_scale=0.08,
distort_color=None):
self._save_path = save_path
train_transforms = self.build_train_transform(distort_color,
resize_scale)
train_dataset = datasets.ImageFolder(self.train_path, train_transforms)
self.local_rank = int(os.environ['LOCAL_RANK'])
self.world_size = int(os.environ['WORLD_SIZE'])
if valid_size is not None:
if isinstance(valid_size, float):
valid_size = int(valid_size * len(train_dataset))
else:
assert isinstance(valid_size,
int), 'invalid valid_size: %s' % valid_size
train_indexes, valid_indexes = self.random_sample_valid_set(
[cls for _, cls in train_dataset.samples],
valid_size,
self.n_classes,
)
train_sampler = torch.utils.data.sampler.SubsetRandomSampler(
train_indexes)
valid_sampler = torch.utils.data.sampler.SubsetRandomSampler(
valid_indexes)
valid_dataset = datasets.ImageFolder(
self.train_path,
transforms.Compose([
transforms.Resize(self.resize_value),
transforms.CenterCrop(self.image_size),
transforms.ToTensor(),
self.normalize,
]))
# self.train = torch.utils.data.DataLoader(
# train_dataset, batch_size=train_batch_size, sampler=train_sampler,
# num_workers=n_worker, pin_memory=True,
# )
from .dali import HybridTrainPipe
from nvidia.dali.plugin.pytorch import DALIClassificationIterator
pipe = HybridTrainPipe(batch_size=train_batch_size,
num_threads=n_worker,
device_id=self.local_rank,
data_dir=self.train_path,
crop=self.image_size)
pipe.build()
self.train = DALIClassificationIterator(
pipe, size=int(pipe.epoch_size("Reader") / self.world_size))
self.valid = torch.utils.data.DataLoader(
valid_dataset,
batch_size=test_batch_size,
sampler=valid_sampler,
num_workers=n_worker,
pin_memory=True,
)
else:
# self.train = torch.utils.data.DataLoader(
# train_dataset, batch_size=train_batch_size, shuffle=True,
# num_workers=n_worker, pin_memory=True,
# )
from .dali import HybridTrainPipe
from nvidia.dali.plugin.pytorch import DALIClassificationIterator
pipe = HybridTrainPipe(batch_size=train_batch_size,
num_threads=n_worker,
device_id=self.local_rank,
data_dir=self.train_path,
crop=self.image_size)
pipe.build()
self.train = DALIClassificationIterator(
pipe, size=int(pipe.epoch_size("Reader") / self.world_size))
self.valid = None
self.test = torch.utils.data.DataLoader(
datasets.ImageFolder(
self.valid_path,
transforms.Compose([
transforms.Resize(self.resize_value),
transforms.CenterCrop(self.image_size),
transforms.ToTensor(),
self.normalize,
])),
batch_size=test_batch_size,
shuffle=False,
num_workers=n_worker,
pin_memory=True,
)
if self.valid is None:
self.valid = self.test
def dataloader(dataset, batch_size, train, workers, length=None, datadir="Data"):
# Dataset
if dataset == "mnist":
mean, std = (0.1307,), (0.3081,)
transform = get_transform(
size=28, padding=0, mean=mean, std=std, preprocess=False
)
dataset = datasets.MNIST(
datadir, train=train, download=True, transform=transform
)
if dataset == "cifar10":
mean, std = (0.491, 0.482, 0.447), (0.247, 0.243, 0.262)
transform = get_transform(
size=32, padding=4, mean=mean, std=std, preprocess=train
)
dataset = datasets.CIFAR10(
datadir, train=train, download=True, transform=transform
)
if dataset == "cifar100":
mean, std = (0.507, 0.487, 0.441), (0.267, 0.256, 0.276)
transform = get_transform(
size=32, padding=4, mean=mean, std=std, preprocess=train
)
dataset = datasets.CIFAR100(
datadir, train=train, download=True, transform=transform
)
if dataset == "tiny-imagenet":
mean, std = (0.480, 0.448, 0.397), (0.276, 0.269, 0.282)
transform = get_transform(
size=64, padding=4, mean=mean, std=std, preprocess=train
)
dataset = custom_datasets.TINYIMAGENET(
datadir, train=train, download=True, transform=transform
)
if dataset == "imagenet":
mean, std = (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
if train:
transform = transforms.Compose(
[
transforms.RandomResizedCrop(224, scale=(0.2, 1.0)),
transforms.RandomGrayscale(p=0.2),
transforms.ColorJitter(0.4, 0.4, 0.4, 0.4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean, std),
]
)
else:
transform = transforms.Compose(
[
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean, std),
]
)
folder = f"{datadir}/imagenet_raw/{'train' if train else 'val'}"
dataset = datasets.ImageFolder(folder, transform=transform)
# Dataloader
use_cuda = torch.cuda.is_available()
kwargs = {"num_workers": workers, "pin_memory": True} if use_cuda else {}
shuffle = train is True
if length is not None:
indices = torch.randperm(len(dataset))[:length]
dataset = torch.utils.data.Subset(dataset, indices)
dataloader = torch.utils.data.DataLoader(
dataset=dataset, batch_size=batch_size, shuffle=shuffle, **kwargs
)
return dataloader
def do_deep_learning(model, data_dir, save_dir, learning_rate, num_epochs,
hidden_units, processor):
if model == 'densenet121':
model = models.densenet121(pretrained=True)
for param in model.parameters():
param.requires_grad = False
from collections import OrderedDict
classifier = nn.Sequential(
OrderedDict([('fc1', nn.Linear(1024, 512)), ('relu1', nn.ReLU()),
('fc2', nn.Linear(512, 2)),
('output', nn.LogSoftmax(dim=1))]))
model.classifier = classifier
elif model == 'alexnet':
model = models.alexnet(pretrained=True)
for param in model.parameters():
param.requires_grad = False
from collections import OrderedDict
classifier = nn.Sequential(
OrderedDict([('fc1', nn.Linear(9216, hidden_units)),
('relu1', nn.ReLU()),
('fc2', nn.Linear(hidden_units, 2)),
('output', nn.LogSoftmax(dim=1))]))
model.classifier = classifier
else:
print("Please select from densenet121 or alexnet only")
train_dir = data_dir + '/train'
valid_dir = data_dir + '/valid'
train_transform = transforms.Compose([
transforms.Resize(255),
transforms.RandomRotation(30),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
#transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
valid_transform = transforms.Compose([
transforms.Resize(255),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
# Load the datasets with ImageFolder
trainset = datasets.ImageFolder(train_dir, transform=train_transform)
validset = datasets.ImageFolder(valid_dir, transform=valid_transform)
# Using the image datasets and the trainforms, define the dataloaders
trainloader = torch.utils.data.DataLoader(trainset,
batch_size=64,
shuffle=True)
validloader = torch.utils.data.DataLoader(validset,
batch_size=32,
shuffle=True)
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.classifier.parameters(), lr=learning_rate)
epochs = num_epochs
print_every = 42
steps = 0
# change to cuda
model.cuda()
for e in range(epochs):
running_loss = 0
for ii, (inputs, labels) in enumerate(trainloader):
steps += 1
inputs, labels = inputs.to('cuda'), labels.to('cuda')
optimizer.zero_grad()
# Forward and backward passes
outputs = model.forward(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
running_loss += loss.item()
if steps % print_every == 0:
model.eval()
with torch.no_grad():
test_loss, accuracy = validation(model, validloader,
criterion)
print(
"Epoch: {}/{}.. ".format(e + 1, epochs),
"Training Loss: {:.3f}.. ".format(running_loss /
print_every),
"Test Loss: {:.3f}.. ".format(test_loss /
len(validloader)),
"Test Accuracy: {:.3f}".format(accuracy /
len(validloader)))
running_loss = 0
# TODO: Save the checkpoint
checkpoint = {
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict,
'classIndex': trainset.class_to_idx,
'epoch': 5,
'hidden': hidden_units,
'model': model
}
torch.save(checkpoint, save_dir)
print('\nModel Saved')
return model
transforms.ToTensor(),
transforms.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225]),
transforms.Normalize([0.5, 0.5, 0.5],
[0.5, 0.5, 0.5])])
data_transforms = transforms.Compose([transforms.RandomRotation(30),
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.RandomVerticalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225]),
transforms.Normalize([0.5, 0.5, 0.5],
[0.5, 0.5, 0.5])])
#datasets
train_dataset = datasets.ImageFolder(train_dir, transform = data_transforms)
valid_dataset = datasets.ImageFolder(valid_dir, transform = data_transforms)
test_dataset = datasets.ImageFolder(test_dir, transform = data_transforms)
#loaders
trainloader = torch.utils.data.DataLoader(train_dataset, batch_size = 64, shuffle = True)
validloader = torch.utils.data.DataLoader(valid_dataset, batch_size = 64, shuffle = True)
testloader = torch.utils.data.DataLoader(test_dataset, batch_size = 64, shuffle = True)
with open('cat_to_name.json', 'r') as f:
cat_to_name = json.load(f)
#creating neural network
class Network(nn.Module):
def train(data_dir=None,
train_dir=None,
batch_size=32,
num_epochs=50,
patience=50,
num_classes=None,
val_dir=None,
model_save_dir=None,
model_save_path=None,
pretrained_model_path=None,
model_type='resnet50',
use_default_pretrained_or_download=True):
if not train_dir:
train_dir = data_dir + '/train'
if not val_dir:
val_dir = data_dir + '/val'
if not num_classes:
num_classes = len(os.listdir(train_dir))
if not model_save_path:
model_save_path = model_save_dir + '/model.pkl'
if pretrained_model_path:
model = torch.load(pretrained_model_path)
else:
model_type = resnet_models[model_type]
model = model_type(pretrained=use_default_pretrained_or_download)
in_features = model.fc.in_features
model.fc = torch.nn.Linear(in_features, num_classes)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
model.to(device)
data_transforms = {
'train':
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
# transforms.ColorJitter(brightness=0.5, contrast=0.5, hue=0.5),
# transforms.RandomRotation(20),
transforms.ToTensor(),
# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
transforms.Normalize([0.5, 0.5, 0.5], [0.229, 0.224, 0.225])
]),
'val':
transforms.Compose([
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
]),
}
image_datasets = {
'train':
datasets.ImageFolder(train_dir, transform=data_transforms['train']),
'val':
datasets.ImageFolder(val_dir, transform=data_transforms['val'])
}
dataloaders = {
x: dataloader.DataLoader(image_datasets[x],
batch_size=batch_size,
shuffle=True,
num_workers=4)
for x in ['train', 'val']
}
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# optimizer=optim.Adam(model.parameters(),lr=0.001)
# scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)
criterion = nn.CrossEntropyLoss()
best_acc = 0.0
bad_epochs = 0
for epoch in range(num_epochs):
log_info = []
for phase in ['train', 'val']:
if phase == 'train':
model.train()
else:
model.eval()
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for inputs, labels in dataloaders[phase]:
inputs = inputs.to(device)
labels = labels.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
_, preds = torch.max(outputs, 1)
loss = criterion(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
running_corrects += torch.sum(preds == labels.data)
# if phase == 'train':
# scheduler.step()
epoch_loss = running_loss / len(dataloaders[phase].dataset)
epoch_acc = running_corrects.double() / len(
dataloaders[phase].dataset)
log_info.append('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
# deep copy the model
if phase == 'val':
if epoch_acc > best_acc:
bad_epochs = 0
best_acc = epoch_acc
torch.save(model, model_save_path)
print('New best model saved to %s' % (model_save_path))
else:
if best_acc >= 0.5:
bad_epochs += 1
if bad_epochs >= patience:
print(
"Val accuracy stopped improving, stop training, best accuracy: %s , model was saved at %s"
% (best_acc, model_save_path))
return
print('\t'.join(log_info))
return ax, image
#-------------------------------
# Hyperparameters Defined
#-------------------------------
batch_size = 128
#input_size = 1000
#hidden_size = 120
#num_classes = 8
num_epochs = 10
#batch_size = 64
learning_rate = 1e-3
#----------------------------------------------------------------------------
all_data = datasets.ImageFolder(root='./natural_images')
train_data_len = int(len(all_data) * 0.8)
valid_data_len = int((len(all_data) - train_data_len) / 2)
test_data_len = int(len(all_data) - train_data_len - valid_data_len)
train_data, val_data, test_data = random_split(
all_data, [train_data_len, valid_data_len, test_data_len])
train_data.dataset.transform = image_transforms['train']
val_data.dataset.transform = image_transforms['val']
test_data.dataset.transform = image_transforms['test']
print(len(train_data), len(val_data), len(test_data))
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=True)
trainiter = iter(train_loader)
def train(args):
## Aim for at least 2 epochs, with a dataset of around 30k images for training.
## Benchmarked a 500x500 style image on Intel Xeon [email protected] v3 CPU. Took 5.91(avg) seconds to train on 1 image.
## 10/10 would not recommend CPU training. 2 epochs on 80k images took 6 days of training.
## An Nvidia GPU speeds it up roughly 49-70x, depending on the model, so go for that.
check_paths(args)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
if args.cuda:
torch.cuda.manual_seed(args.seed)
kwargs = {'num_workers': 0, 'pin_memory': False}
else:
kwargs = {}
transform = transforms.Compose([
transforms.Scale(args.image_size),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))
])
train_dataset = datasets.ImageFolder(args.dataset, transform)
train_loader = DataLoader(train_dataset,
batch_size=args.batch_size,
**kwargs)
style_model = Net(ngf=args.ngf)
if args.resume is not None:
print('Resuming, initializing using weight from {}.'.format(
args.resume))
style_model.load_state_dict(torch.load(args.resume))
print(style_model)
optimizer = Adam(style_model.parameters(), args.lr)
mse_loss = torch.nn.MSELoss()
vgg = Vgg16()
utils.init_vgg16(args.vgg_model_dir)
vgg.load_state_dict(
torch.load(os.path.join(args.vgg_model_dir, "vgg16.weight")))
if args.cuda:
style_model.cuda()
vgg.cuda()
style_loader = utils.StyleLoader(args.style_folder, args.style_size)
tbar = trange(args.epochs)
for e in tbar:
style_model.train()
agg_content_loss = 0.
agg_style_loss = 0.
count = 0
for batch_id, (x, _) in enumerate(train_loader):
n_batch = len(x)
count += n_batch
optimizer.zero_grad()
x = Variable(utils.preprocess_batch(x))
if args.cuda:
x = x.cuda()
style_v = style_loader.get(batch_id)
style_model.setTarget(style_v)
style_v = utils.subtract_imagenet_mean_batch(style_v)
features_style = vgg(style_v)
gram_style = [utils.gram_matrix(y) for y in features_style]
y = style_model(x)
xc = Variable(x.data.clone())
y = utils.subtract_imagenet_mean_batch(y)
xc = utils.subtract_imagenet_mean_batch(xc)
features_y = vgg(y)
features_xc = vgg(xc)
f_xc_c = Variable(features_xc[1].data, requires_grad=False)
content_loss = args.content_weight * mse_loss(
features_y[1], f_xc_c)
style_loss = 0.
for m in range(len(features_y)):
gram_y = utils.gram_matrix(features_y[m])
gram_s = Variable(gram_style[m].data,
requires_grad=False).repeat(
args.batch_size, 1, 1, 1)
style_loss += args.style_weight * mse_loss(
gram_y, gram_s[:n_batch, :, :])
total_loss = content_loss + style_loss
total_loss.backward()
optimizer.step()
agg_content_loss += content_loss.data[0]
agg_style_loss += style_loss.data[0]
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.6f}\tstyle: {:.6f}\ttotal: {:.6f}".format(
time.ctime(), e + 1, count, len(train_dataset),
agg_content_loss / (batch_id + 1),
agg_style_loss / (batch_id + 1),
(agg_content_loss + agg_style_loss) / (batch_id + 1))
tbar.set_description(mesg)
if (batch_id + 1) % (4 * args.log_interval) == 0:
# save model
style_model.eval()
style_model.cpu()
save_model_filename = "Epoch_" + str(e) + "iters_" + str(count) + "_" + \
str(time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir,
save_model_filename)
torch.save(style_model.state_dict(), save_model_path)
style_model.train()
style_model.cuda()
tbar.set_description("\nCheckpoint, trained model saved at",
save_model_path)
# save model
style_model.eval()
style_model.cpu()
save_model_filename = "Final_epoch_" + str(args.epochs) + "_" + \
str(time.ctime()).replace(' ', '_') + "_" + str(
args.content_weight) + "_" + str(args.style_weight) + ".model"
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
torch.save(style_model.state_dict(), save_model_path)
print("\nDone, trained model saved at", save_model_path)
def train(args):
device = torch.device("cuda" if args.cuda else "cpu")
np.random.seed(args.seed)
torch.manual_seed(args.seed)
transform = transforms.Compose([
transforms.Resize(args.image_size),
transforms.CenterCrop(args.image_size),
transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255)),
])
train_dataset = datasets.ImageFolder(args.dataset, transform)
train_loader = DataLoader(train_dataset, batch_size=args.batch_size)
transformer = TransformerNet().to(device)
optimizer = Adam(transformer.parameters(), args.lr)
mse_loss = torch.nn.MSELoss()
vgg = Vgg16(requires_grad=False).to(device)
style_transform = transforms.Compose(
[transforms.ToTensor(),
transforms.Lambda(lambda x: x.mul(255))])
style = utils.load_image(args.style_image, size=args.style_size)
style = style_transform(style)
style = style.repeat(args.batch_size, 1, 1, 1).to(device)
features_style = vgg(utils.normalize_batch(style))
gram_style = [utils.gram_matrix(y) for y in features_style]
for e in range(args.epochs):
transformer.train()
agg_content_loss = 0.0
agg_style_loss = 0.0
count = 0
for batch_id, (x, _) in enumerate(train_loader):
n_batch = len(x)
count += n_batch
optimizer.zero_grad()
x = x.to(device)
y = transformer(x)
y = utils.normalize_batch(y)
x = utils.normalize_batch(x)
features_y = vgg(y)
features_x = vgg(x)
content_loss = args.content_weight * mse_loss(
features_y.relu2_2, features_x.relu2_2)
style_loss = 0.0
for ft_y, gm_s in zip(features_y, gram_style):
gm_y = utils.gram_matrix(ft_y)
style_loss += mse_loss(gm_y, gm_s[:n_batch, :, :])
style_loss *= args.style_weight
total_loss = content_loss + style_loss
total_loss.backward()
optimizer.step()
agg_content_loss += content_loss.item()
agg_style_loss += style_loss.item()
if (batch_id + 1) % args.log_interval == 0:
mesg = "{}\tEpoch {}:\t[{}/{}]\tcontent: {:.6f}\tstyle: {:.6f}\ttotal: {:.6f}".format(
time.ctime(),
e + 1,
count,
len(train_dataset),
agg_content_loss / (batch_id + 1),
agg_style_loss / (batch_id + 1),
(agg_content_loss + agg_style_loss) / (batch_id + 1),
)
print(mesg)
if (args.checkpoint_model_dir is not None
and (batch_id + 1) % args.checkpoint_interval == 0):
transformer.eval().cpu()
ckpt_model_filename = ("ckpt_epoch_" + str(e) + "_batch_id_" +
str(batch_id + 1) + ".pth")
ckpt_model_path = os.path.join(args.checkpoint_model_dir,
ckpt_model_filename)
torch.save(transformer.state_dict(), ckpt_model_path)
transformer.to(device).train()
# save model
transformer.eval().cpu()
save_model_filename = ("epoch_" + str(args.epochs) + "_" +
str(time.ctime()).replace(" ", "_") + "_" +
str(args.content_weight) + "_" +
str(args.style_weight) + ".model")
save_model_path = os.path.join(args.save_model_dir, save_model_filename)
torch.save(transformer.state_dict(), save_model_path)
print("\nDone, trained model saved at", save_model_path)
train_data = dset.CIFAR100('/share/data/vision-greg/cifarpy',
train=True,
transform=train_transform,
download=False)
test_data = dset.CIFAR100('/share/data/vision-greg/cifarpy',
train=False,
transform=test_transform,
download=False)
num_classes = 100
else:
train_data = dset.ImageFolder(
'/share/data/vision-greg/DistortedImageNet/Icons-50',
transform=trn.Compose([
trn.Resize((32, 32)),
trn.RandomHorizontalFlip(),
trn.RandomCrop(32, padding=4),
trn.ToTensor(),
# RandomErasing()
]))
test_data = dset.ImageFolder(
'/share/data/vision-greg/DistortedImageNet/Icons-50',
transform=trn.Compose([trn.Resize((32, 32)),
trn.ToTensor()]))
num_classes = 50
if args.traditional:
filtered_imgs = []
for img in train_data.samples:
img_name = img[0]
if '_2' not in img_name:
metavar='ARCH',
default='proxyless_mobile_14',
choices=model_names,
help='model architecture: ' + ' | '.join(model_names) +
' (default: proxyless_mobile_14)')
parser.add_argument('--manual_seed', default=0, type=int)
args = parser.parse_args()
net = tf_model_zoo.__dict__[args.arch](pretrained=True)
data_loader = torch.utils.data.DataLoader(
datasets.ImageFolder(
osp.join(args.path, "val"),
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]),
])),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
pin_memory=True,
drop_last=False,
)
losses = AverageMeter()
top1 = AverageMeter()
for i, (_input, target) in enumerate(data_loader):
images = _input.numpy()
# Horovod: print logs on the first worker.
verbose = 1 if hvd.rank() == 0 else 0
# Horovod: write TensorBoard logs on first worker.
log_writer = tensorboardX.SummaryWriter(
args.log_dir) if hvd.rank() == 0 else None
# Horovod: limit # of CPU threads to be used per worker.
torch.set_num_threads(4)
kwargs = {'num_workers': 4, 'pin_memory': True} if args.cuda else {}
train_dataset = \
datasets.ImageFolder(args.train_dir,
transform=transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]))
# Horovod: use DistributedSampler to partition data among workers. Manually specify
# `num_replicas=hvd.size()` and `rank=hvd.rank()`.
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset, num_replicas=hvd.size(), rank=hvd.rank())
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=allreduce_batch_size,
sampler=train_sampler,
**kwargs)
val_dataset = \
datasets.ImageFolder(args.val_dir,
transform=transforms.Compose([
if isCrop:
transform = transforms.Compose([
transforms.Scale(108),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
else:
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
dset = datasets.ImageFolder(data_dir, transform)
train_loader = torch.utils.data.DataLoader(dset,
batch_size=batch_size,
shuffle=True)
temp = plt.imread(train_loader.dataset.imgs[0][0])
if (temp.shape[0] != img_size) or (temp.shape[0] != img_size):
sys.stderr.write(
'Error! image size is not 64 x 64! run \"celebA_data_preprocess.py\" !!!'
)
sys.exit(1)
dset = datasets.ImageFolder(personal_dir, transform)
personal_loader = torch.utils.data.DataLoader(dset,
batch_size=3,
shuffle=False)
def main():
cudnn.benchmark = True
best_prec1 = 0
if args.deterministic:
cudnn.benchmark = False
cudnn.deterministic = True
torch.manual_seed(args.local_rank)
torch.set_printoptions(precision=10)
# handle distributed traininc
args.distributed = False
if 'WORLD_SIZE' in os.environ:
args.distributed = int(os.environ['WORLD_SIZE']) > 1
args.gpu = 0
args.world_size = 1
if args.distributed:
args.gpu = args.local_rank
torch.cuda.set_device(args.gpu)
torch.distributed.init_process_group(backend='nccl',
init_method='env://')
args.world_size = torch.distributed.get_world_size()
assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."
global base_learning_rate
global best_acc1
base_learning_rate = args.base_lr * float(
args.batch_size * args.world_size) / 256.
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
model = model.cuda()
# define loss function (criterion) and optimizer
cel = nn.CrossEntropyLoss()
criterion = lambda pred, target, lam: (-F.log_softmax(
pred, dim=1) * torch.zeros(pred.size()).cuda().scatter_(
1, target.data.view(-1, 1), lam.view(-1, 1))).sum(dim=1).mean()
parameters_bias = [
p[1] for p in model.named_parameters() if 'bias' in p[0]
]
parameters_scale = [
p[1] for p in model.named_parameters() if 'scale' in p[0]
]
parameters_others = [
p[1] for p in model.named_parameters()
if not ('bias' in p[0] or 'scale' in p[0])
]
optimizer = torch.optim.SGD([{
'params': parameters_bias,
'lr': args.base_lr / 10.
}, {
'params': parameters_scale,
'lr': args.base_lr / 10.
}, {
'params': parameters_others
}],
lr=base_learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
# Initialize Amp. Amp accepts either values or strings for the optional override arguments,
# for convenient interoperation with argparse.
model, optimizer = amp.initialize(
model,
optimizer,
opt_level=args.opt_level,
keep_batchnorm_fp32=args.keep_batchnorm_fp32,
loss_scale=args.loss_scale)
#torch.cuda.set_device(args.gpu)
model = DDP(model, delay_allreduce=True)
#model = torch.nn.DataParallel(model)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion, args)
return
sgdr = CosineAnnealingLR(optimizer, args.epochs, eta_min=0, last_epoch=-1)
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
# evaluate on validation set
acc1 = validate(val_loader, model, cel, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if args.local_rank == 0:
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)
#X = torch.load(ID)
img = Image.open(ID) # use pillow to open a file
img = img.convert('RGB') #convert image to RGB channel
if self.transform is not None:
img = self.transform(img)
#img = np.asarray(img).transpose(-1, 0, 1) # we have to change the dimensions from width x height x channel (WHC) to channel x width x height (CWH)
img = torch.from_numpy(np.asarray(img)) # create the image tensor
X = img
y = self.labels[ID]
return X, y
image_datasets = {
x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
for x in ['train', 'val']
}
partition = {}
partition['train'] = []
partition['val'] = []
partition['val_photo'] = []
labels2 = {}
base_mapping = image_datasets['train'].class_to_idx
print("Preprocessing datasets")
class_name_to_id = image_datasets['train'].class_to_idx
def main():
global args, best_prec1
args = parser.parse_args()
if args.seed is not None:
random.seed(args.seed)
torch.manual_seed(args.seed)
cudnn.deterministic = True
warnings.warn('You have chosen to seed training. '
'This will turn on the CUDNN deterministic setting, '
'which can slow down your training considerably! '
'You may see unexpected behavior when restarting '
'from checkpoints.')
if args.gpu is not None:
warnings.warn('You have chosen a specific GPU. This will completely '
'disable data parallelism.')
args.distributed = args.world_size > 1
if args.distributed:
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if args.gpu is not None:
model = model.cuda(args.gpu)
elif args.distributed:
model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model)
else:
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss().cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
args.resume = tmp_resume
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = True
# Data loading code
traindir = os.path.join(args.tmp_data_dir, 'train')
valdir = os.path.join(args.tmp_data_dir, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion)
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch)
# evaluate on validation set
prec1 = validate(val_loader, model, criterion)
# remember best prec@1 and save checkpoint
is_best = prec1 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
}, is_best)
# timely backup
source = os.path.join(args.tmp_save_dir)
target = os.path.join(args.train_url)
mox.file.copy_parallel(source, target)
transform = transforms.Compose([
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
else:
transform = transforms.Compose([
transforms.Resize(64),
transforms.CenterCrop(64),
transforms.Grayscale(num_output_channels=1),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5]),
])
# Create the dataset
dataset = dset.ImageFolder(root=dataroot, transform=transform)
# Create the dataloader
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers)
# custom weights initialization called on netG and netD
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find('BatchNorm') != -1:
nn.init.normal_(m.weight.data, 1.0, 0.02)
nn.init.constant_(m.bias.data, 0)
transforms.ColorJitter(
brightness=0.1, contrast=0.1, saturation=0.1, hue=0)
] + transform_train_list
print(transform_train_list)
data_transforms = {
'train': transforms.Compose(transform_train_list),
'val': transforms.Compose(transform_val_list),
}
train_all = ''
if opt.train_all:
train_all = '_all'
image_datasets = {}
image_datasets['train'] = datasets.ImageFolder(
os.path.join(data_dir, 'train' + train_all), data_transforms['train'])
image_datasets['val'] = datasets.ImageFolder(os.path.join(data_dir, 'val'),
data_transforms['val'])
dataloaders = {
x: torch.utils.data.DataLoader(image_datasets[x],
batch_size=opt.batchsize,
shuffle=True,
num_workers=16)
for x in ['train', 'val']
}
dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes
use_gpu = torch.cuda.is_available()
def main_worker(gpu, ngpus_per_node, args):
global best_acc1
print("Use CPU: {} for training".format(gpu))
if args.distributed:
if args.dist_url == "env://" and args.rank == -1:
args.rank = int(os.environ["RANK"])
if args.multiprocessing_distributed:
# For multiprocessing distributed training, rank needs to be the
# global rank among all the processes
args.rank = args.rank * ngpus_per_node + gpu
dist.init_process_group(backend=args.dist_backend,
init_method=args.dist_url,
world_size=args.world_size,
rank=args.rank)
# create model
if args.pretrained:
print("=> using pre-trained model '{}'".format(args.arch))
model = models.__dict__[args.arch](pretrained=True, quantize=False)
else:
print("=> creating model '{}'".format(args.arch))
model = models.__dict__[args.arch]()
if args.distributed:
# For multiprocessing distributed, DistributedDataParallel constructor
# should always set the single device scope, otherwise,
# DistributedDataParallel will use all available devices.
if args.gpu is not None:
torch.cuda.set_device(args.gpu)
model.cuda(args.gpu)
# When using a single GPU per process and per
# DistributedDataParallel, we need to divide the batch size
# ourselves based on the total number of GPUs we have
args.batch_size = int(args.batch_size / ngpus_per_node)
args.workers = int(args.workers / ngpus_per_node)
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
else:
model = torch.nn.parallel.DistributedDataParallelCPU(model)
elif args.gpu is not None:
torch.cuda.set_device(args.gpu)
model = model.cuda(args.gpu)
else:
# DataParallel will divide and allocate batch_size to all available GPUs
if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
model.features = torch.nn.DataParallel(model.features)
model.cuda()
else:
model = torch.nn.DataParallel(model)
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_acc1 = checkpoint['best_acc1']
if args.gpu is not None:
# best_acc1 may be from a checkpoint from a different GPU
best_acc1 = best_acc1.to(args.gpu)
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
# Data loading code
traindir = os.path.join(args.data, 'train')
valdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_dataset = datasets.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
normalize,
]))
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
train_dataset)
else:
train_sampler = None
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler)
val_loader = torch.utils.data.DataLoader(datasets.ImageFolder(
valdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
if args.evaluate:
validate(val_loader, model, criterion, args)
if args.tune:
model.eval()
model.module.fuse_model()
from lpot import Quantization, common
quantizer = Quantization("./conf_dump_tensors.yaml")
quantizer.model = common.Model(model)
q_model = quantizer()
return
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
adjust_learning_rate(optimizer, epoch, args)
# train for one epoch
train(train_loader, model, criterion, optimizer, epoch, args)
# evaluate on validation set
acc1 = validate(val_loader, model, criterion, args)
# remember best acc@1 and save checkpoint
is_best = acc1 > best_acc1
best_acc1 = max(acc1, best_acc1)
if not args.multiprocessing_distributed or (
args.multiprocessing_distributed
and args.rank % ngpus_per_node == 0):
save_checkpoint(
{
'epoch': epoch + 1,
'arch': args.arch,
'state_dict': model.state_dict(),
'best_acc1': best_acc1,
'optimizer': optimizer.state_dict(),
}, is_best)
def main():
if not torch.cuda.is_available():
logging.info('no gpu device available')
sys.exit(1)
np.random.seed(args.seed)
torch.cuda.set_device(args.gpu)
cudnn.benchmark = True
torch.manual_seed(args.seed)
cudnn.enabled = True
torch.cuda.manual_seed(args.seed)
logging.info('gpu device = %d' % args.gpu)
logging.info("args = %s", args)
genotype = eval("genotypes.%s" % args.arch)
model = Network(args.init_channels, CLASSES, args.layers, args.auxiliary,
genotype)
if args.parallel:
model = nn.DataParallel(model).cuda()
else:
model = model.cuda()
logging.info("param size = %fMB", utils.count_parameters_in_MB(model))
criterion = nn.CrossEntropyLoss()
criterion = criterion.cuda()
criterion_smooth = CrossEntropyLabelSmooth(CLASSES, args.label_smooth)
criterion_smooth = criterion_smooth.cuda()
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
traindir = os.path.join(args.data, 'train')
validdir = os.path.join(args.data, 'val')
normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
train_data = dset.ImageFolder(
traindir,
transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ColorJitter(brightness=0.4,
contrast=0.4,
saturation=0.4,
hue=0.2),
transforms.ToTensor(),
normalize,
]))
valid_data = dset.ImageFolder(
validdir,
transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
normalize,
]))
train_queue = torch.utils.data.DataLoader(train_data,
batch_size=args.batch_size,
shuffle=True,
pin_memory=True,
num_workers=4)
valid_queue = torch.utils.data.DataLoader(valid_data,
batch_size=64,
shuffle=False,
pin_memory=True,
num_workers=4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
args.decay_period,
gamma=args.gamma)
best_acc_top1 = 0
for epoch in range(args.epochs):
scheduler.step()
logging.info('epoch %d lr %e', epoch, scheduler.get_lr()[0])
model.drop_path_prob = args.drop_path_prob * epoch / args.epochs
train_acc, train_obj = train(train_queue, model, criterion_smooth,
optimizer)
logging.info('train_acc %f', train_acc)
valid_acc_top1, valid_acc_top5, valid_obj = infer(
valid_queue, model, criterion)
logging.info('valid_acc_top1 %f', valid_acc_top1)
logging.info('valid_acc_top5 %f', valid_acc_top5)
is_best = False
if valid_acc_top1 > best_acc_top1:
best_acc_top1 = valid_acc_top1
is_best = True
utils.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_acc_top1': best_acc_top1,
'optimizer': optimizer.state_dict(),
}, is_best, args.save)
train_loader = transforms.Compose([
transforms.RandomResizedCrop(224),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
test_loader = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
data_dir = 'hymenoptera_data'
train_image_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), train_loader)
test_image_dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'), test_loader)
train_image_dataloader = torch.utils.data.DataLoader(train_image_dataset,
batch_size=4, shuffle=True, num_workers=0)
test_image_dataloader = torch.utils.data.DataLoader(test_image_dataset,
batch_size=4, shuffle=True, num_workers=0)
dataset_sizes = {'train': len(train_image_dataset), 'val': len(test_image_dataset)}
class_names = train_image_dataset.classes
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
def imshow(inp, title=None):
inp = inp.numpy().transpose((1,2,0)) # inp = [3, 228, 906]. after transpose inp = [228, 906, 3]
mean = np.array([0.485, 0.456, 0.406])
std = np.array([0.229, 0.224, 0.225])
if isinstance(custom_size, list): #判断custom_size是否为list
img_size = custom_size
tmp = "Image size used:\t{0}x{1}".format(img_size[0], img_size[1])
utils.print_both(f, tmp)
# Transformations
data_transforms = transforms.Compose([
transforms.Resize(img_size[0:2]),
# transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) #对额外加入的数据进行处理
])
# Read data from selected folder and apply transformations
image_dataset = datasets.ImageFolder(data_dir, data_transforms)
# Prepare data for network: schuffle and arrange batches
dataloader = torch.utils.data.DataLoader(image_dataset, batch_size=batch,
shuffle=False, num_workers=workers)
# Size of data sets
dataset_size = len(image_dataset)
tmp = "Training set size:\t" + str(dataset_size)
utils.print_both(f, tmp)
params['dataset_size'] = dataset_size
# GPU check
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") #torch.device包含一个设备类型（‘cpu’或‘cuda’）和可选的设备序号。
tmp = "\nPerforming calculations on:\t" + str(device)
utils.print_both(f, tmp + '\n')
