def __init__(self, args):
self.opt = args
self.data_path = self.opt.data_path
self.model_path = self.opt.model_path
self.model_name = self.opt.model_name
self.batch_size = self.opt.batch_size
self.num_workers = self.opt.num_workers
self.input_size = (self.opt.height, self.opt.width)
# Load model for testing
model, optimizer, criterion, epoch, device = load_model(os.path.join(self.model_path, self.model_name))
model.eval()
self.model = model
self.device = device
self.criterion = criterion
self.transforms = transforms.Compose([transforms.Resize(self.input_size),
transforms.ToTensor()])
# List of perturbations we'll test the model on
self.perturbations = ['none', 'fog', 'rain', 'snow', 'occlusion']
# Dataset and dataloader dictionaries indexed by the type of perturbation it applies to images
self.datasets = {i: Cityscapes(self.data_path,
split='val',
mode='coarse',
target_type=['polygon'],
transform=self.transforms,
target_transform=get_image_labels,
perturbation=i)
for i in self.perturbations}
self.num_testing_files = self.datasets['none'].__len__()
print('--> Testing model {} on {} test files\n'.format(self.model_name, self.num_testing_files))
# Creating PyTorch dataloaders
self.dataloaders = {i: DataLoader(self.datasets[i], batch_size=self.batch_size, shuffle=True,
num_workers=self.num_workers) for i in self.perturbations}
parser.add_argument('-summary', '--summary', action='store_true')
args = parser.parse_args()
all_dev = parse_devices(args.devices)
network = BiSeNet(config.num_classes,
is_training=False,
criterion=None,
ohem_criterion=None)
data_setting = {
'img_root': config.img_root_folder,
'gt_root': config.gt_root_folder,
'train_source': config.train_source,
'eval_source': config.eval_source
}
dataset = Cityscapes(data_setting, 'val', None)
if args.speed_test:
device = all_dev[0]
logger.info("=========DEVICE:%s SIZE:%s=========" %
(torch.cuda.get_device_name(device), args.input_size))
input_size = tuple(int(x) for x in args.input_size.split('x'))
compute_speed(network, input_size, device, args.iteration)
elif args.summary:
input_size = tuple(int(x) for x in args.input_size.split('x'))
stat(network, input_size)
else:
with torch.no_grad():
segmentor = SegEvaluator(dataset, config.num_classes,
config.image_mean, config.image_std,
network, config.eval_scale_array,
def main_worker(gpu, ngpus_per_node, args):
global best_mIoU
args.gpu = gpu
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 'DFANet'")
model = DFANet(pretrained=True, pretrained_backbone=False)
else:
print("=> creating model 'DFANet'")
model = DFANet(pretrained=False, pretrained_backbone=True)
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.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
model = torch.nn.DataParallel(model).cuda()
# define loss function (criterion) and optimizer
criterion = nn.CrossEntropyLoss(ignore_index=19).cuda(args.gpu)
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
metric = IoU(20, ignore_index=19)
# 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_mIoU = checkpoint['best_mIoU']
if args.gpu is not None:
# best_mIoU may be from a checkpoint from a different GPU
best_mIoU = best_mIoU.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
train_dataset = Cityscapes(args.data,
split='train',
mode='fine',
target_type='semantic',
transform=joint_transforms.Compose([
joint_transforms.RandomHorizontalFlip(),
joint_transforms.RandomSized(1024),
joint_transforms.ToTensor(),
joint_transforms.Normalize(
mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]))
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(Cityscapes(
args.data,
split='val',
mode='fine',
target_type='semantic',
transform=joint_transforms.Compose([
joint_transforms.RandomHorizontalFlip(),
joint_transforms.RandomSized(1024),
joint_transforms.ToTensor(),
joint_transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers,
pin_memory=True)
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)
# evaluate on training data
train_mIoU, train_loss = validate(train_loader, model, criterion,
metric, args)
# evaluate on validation set
val_mIoU, val_loss = validate(val_loader, model, criterion, metric,
args)
print("Train mIoU: {}".format(train_mIoU))
print("Train Loss: {}".format(train_loss))
print("Val mIoU: {}".format(val_mIoU))
print("Val mIoU: {}".format(val_loss))
crop_size = 520
train_transform = et.ExtCompose([
et.ExtRandomCrop(size=(crop_size, crop_size)),
et.ExtColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
et.ExtRandomHorizontalFlip(),
et.ExtToTensor(),
et.ExtNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
val_transform = et.ExtCompose([
et.ExtToTensor(),
et.ExtNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
train_dst = Cityscapes(root=data_root,
split='train',
transform=train_transform)
val_dst = Cityscapes(root=data_root, split='val', transform=val_transform)
train_loader = data.DataLoader(train_dst,
batch_size=2,
shuffle=True,
num_workers=2)
val_loader = data.DataLoader(val_dst,
batch_size=2,
shuffle=True,
num_workers=2)
for i, data in enumerate(train_loader):
images, depths, labels = data
if i == 0:
def __init__(self, args):
self.opt = args
self.data_path = self.opt.data_path
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.model_path = self.opt.model_path
self.model_name = self.opt.model_name
self.model_type = self.opt.model_type
self.pretrained = self.opt.pretrained
self.freeze = self.opt.freeze
self.finetune = self.opt.finetune
self.train_on_perturbations = self.opt.perturbations
# Training parameters
self.input_size = (self.opt.height, self.opt.width)
self.batch_size = int(self.opt.batch_size)
self.num_workers = int(self.opt.num_workers)
self.epochs = int(self.opt.num_epochs)
self.lr = float(self.opt.learning_rate)
self.step = int(self.opt.scheduler_step_size)
# Create model and place on GPU
assert self.model_type in ['inception', 'resnet', 'VGG']
if self.model_type == 'inception':
self.model = inception_v3(pretrained=self.pretrained,
aux_logits=False)
if not self.finetune:
set_requires_grad(self.model, requires_grad=(not self.freeze))
else:
layers_to_train = ['Mixed_7a', 'Mixed_7b', 'Mixed_7c', 'fc']
set_requires_grad(self.model,
finetune=self.finetune,
layers_to_train=layers_to_train)
self.model.fc = self.fc = nn.Linear(2048, len(important_classes))
elif self.model_type == 'VGG':
self.model = vgg16_bn(pretrained=self.pretrained)
if not self.finetune:
set_requires_grad(self.model, requires_grad=(not self.freeze))
else:
layers_to_train = [
'features34', 'features35', 'features37', 'features38',
'features40', 'features41', 'classifier'
]
set_requires_grad(self.model,
finetune=self.finetune,
layers_to_train=layers_to_train,
vgg=True)
num_ftrs = self.model.classifier[6].in_features
self.model.classifier[6] = nn.Linear(num_ftrs,
len(important_classes))
else:
self.model = resnet34(pretrained=self.pretrained)
if not self.finetune:
set_requires_grad(self.model, requires_grad=(not self.freeze))
else:
layers_to_train = ['layer4', 'fc']
set_requires_grad(self.model,
finetune=self.finetune,
layers_to_train=layers_to_train)
self.model.fc = self.fc = nn.Linear(512, len(important_classes))
self.model.to(self.device)
# Loss function and optimizer
self.criterion = nn.BCEWithLogitsLoss()
params_to_update = []
for name, param in self.model.named_parameters():
if param.requires_grad:
params_to_update.append(param)
print('Training param {}'.format(name))
# self.optimizer = optim.Adam(params_to_update, lr=self.lr)
self.optimizer = optim.SGD(params_to_update,
lr=self.lr,
momentum=0.9,
weight_decay=0.0005)
self.scheduler = StepLR(self.optimizer, step_size=self.step, gamma=0.9)
print(
'Training options:\n'
'\tModel: {}\n\Pretrained: {}\n\tInput size: {}\n\tBatch size: {}\n'
'\tEpochs: {}\n\t'
'Learning rate: {}\n\tStep Size: {}\n'.format(
self.model_type.capitalize(), self.pretrained, self.input_size,
self.batch_size, self.epochs, self.lr, self.step))
# Data transformations to be used during loading of images
self.data_transforms = {
'train':
transforms.Compose(
[transforms.Resize(self.input_size),
transforms.ToTensor()]),
'val':
transforms.Compose(
[transforms.Resize(self.input_size),
transforms.ToTensor()]),
'test':
transforms.Compose(
[transforms.Resize(self.input_size),
transforms.ToTensor()])
}
perturbation = 'none'
if self.train_on_perturbations:
perturbation = 'random'
# Creating PyTorch datasets
self.datasets = dict()
train_dataset = Cityscapes(self.data_path,
split='train',
mode='fine',
target_type=["polygon"],
transform=self.data_transforms['train'],
target_transform=get_image_labels,
perturbation=perturbation)
trainextra_dataset = Cityscapes(
self.data_path,
split='train_extra',
mode='coarse',
target_type=["polygon"],
transform=self.data_transforms['train'],
target_transform=get_image_labels,
perturbation=perturbation)
self.datasets['train'] = ConcatDataset(
[train_dataset, trainextra_dataset])
self.datasets['val'] = Cityscapes(
self.data_path,
split='val',
mode='coarse',
target_type=['polygon'],
transform=self.data_transforms['val'],
target_transform=get_image_labels)
self.datasets['test'] = Cityscapes(
self.data_path,
split='test',
mode='fine',
target_type=["polygon"],
transform=self.data_transforms['test'],
target_transform=get_image_labels)
self.dataset_lens = [
self.datasets[i].__len__() for i in ['train', 'val', 'test']
]
print('Training on:\n'
'\tTrain files: {}\n\tValidation files: {}\n\tTest files: {}\n' \
.format(*self.dataset_lens))
# Creating PyTorch dataloaders
self.dataloaders = {
i: DataLoader(self.datasets[i],
batch_size=self.batch_size,
shuffle=True,
num_workers=self.num_workers)
for i in ['train', 'val']
}
self.test_loader = DataLoader(dataset=self.datasets['test'],
batch_size=1,
shuffle=True)
def test_filename_conversion(self):
img_filename = "datasets/cityscapes/leftImg8bit/train/aachen/aachen_000000_000019_leftImg8bit.png"
label_filename = "datasets/cityscapes/gtFine/train/aachen/aachen_000000_000019_gtFine_labelIds.png"
self.assertEqual(Cityscapes.img_to_label_filename(img_filename), label_filename)