def __init__(self, args, student_model, teacher_model, src_loader,
trg_loader, val_loader, optimizer, teacher_optimizer):
self.args = args
self.student_model = student_model
self.teacher_model = teacher_model
self.src_loader = src_loader
self.trg_loader = trg_loader
self.val_loader = val_loader
self.optimizer = optimizer
self.teacher_optimizer = teacher_optimizer
# Define Evaluator
self.evaluator = Evaluator(args.nclass)
# Define lr scheduler
# self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
# args.epochs, len(trn_loader))
#self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer, milestones=[3, 6, 9, 12], gamma=0.5)
#ft
self.scheduler = torch.optim.lr_scheduler.MultiStepLR(self.optimizer,
milestones=[20],
gamma=0.5)
self.best_pred = 0
self.init_weight = 0.98
# Define Saver
self.saver = Saver(self.args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.evaluator = Evaluator(self.args.nclass)
def setup_saver_and_summary(self, num_current_labeled_samples, samples, experiment_group=None, regions=None):
self.saver = ActiveSaver(self.args, num_current_labeled_samples, experiment_group=experiment_group)
self.saver.save_experiment_config()
self.saver.save_active_selections(samples, regions)
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.num_current_labeled_samples = num_current_labeled_samples
def __init__(self, args):
self.args = args
# define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
# Define weight
self.temporal_weight = args.temporal_weight
self.spatial_weight = args.spatial_weight
# Define network
temporal_model = Model(name='vgg16_bn', num_classes=101,
is_flow=True).get_model()
spatial_model = Model(name='vgg16_bn', num_classes=101,
is_flow=False).get_model()
# Define Optimizer
#optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
temporal_optimizer = torch.optim.Adam(temporal_model.parameters(),
lr=args.temporal_lr)
spatial_optimizer = torch.optim.Adam(spatial_model.parameters(),
lr=args.spatial_lr)
# Define Criterion
self.temporal_criterion = nn.BCELoss().cuda()
self.spatial_criterion = nn.BCELoss().cuda()
self.temporal_model, self.temporal_optimizer = temporal_model, temporal_optimizer
self.spatial_model, self.spatial_optimizer = spatial_model, spatial_optimizer
# Define Evaluator
self.top1_eval = Evaluator(self.nclass)
# Using cuda
if args.cuda:
self.temporal_model = torch.nn.DataParallel(
self.temporal_model, device_ids=self.args.gpu_ids)
patch_replication_callback(self.temporal_model)
self.temporal_model = self.temporal_model.cuda()
self.spatial_model = torch.nn.DataParallel(
self.spatial_model, device_ids=self.args.gpu_ids)
patch_replication_callback(self.spatial_model)
self.spatial_model = self.spatial_model.cuda()
# Resuming checkpoint
self.best_accuracy = 0.0
'''
def __init__(self, config, args):
self.args = args
self.config = config
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader = make_data_loader(
config)
# Define network
#self.model = DeepLab(num_classes=self.nclass,
# backbone=config.backbone,
# output_stride=config.out_stride,
# sync_bn=config.sync_bn,
# freeze_bn=config.freeze_bn)
self.model = UNet(n_channels=1, n_classes=3, bilinear=True)
#train_params = [{'params': self.model.get_1x_lr_params(), 'lr': config.lr},
# {'params': self.model.get_10x_lr_params(), 'lr': config.lr * config.lr_ratio}]
# Define Optimizer
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=config.lr,
momentum=config.momentum,
weight_decay=config.weight_decay)
# Define Criterion
# whether to use class balanced weights
self.criterion = MSELoss(cuda=args.cuda)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.best_pred_source = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def __init__(self, args):
self.args = args
self.saver = Saver(args)
self.saver.save_experiment_config()
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.logger = self.saver.create_logger()
kwargs = {'num_workers': args.workers, 'pin_memory': False}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args, **kwargs)
self.model = EDCNet(args.rgb_dim, args.event_dim, num_classes=self.nclass, use_bn=True)
train_params = [{'params': self.model.random_init_params(),
'lr': 10*args.lr, 'weight_decay': 10*args.weight_decay},
{'params': self.model.fine_tune_params(),
'lr': args.lr, 'weight_decay': args.weight_decay}]
self.optimizer = torch.optim.Adam(train_params, lr=args.lr, weight_decay=args.weight_decay)
if args.cuda:
self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.to(self.args.device)
if args.use_balanced_weights:
root_dir = Path.db_root_dir(args.dataset)[0] if isinstance(Path.db_root_dir(args.dataset), list) else Path.db_root_dir(args.dataset)
classes_weights_path = os.path.join(root_dir,
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass, classes_weights_path)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.criterion_event = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode='event')
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs, len(self.train_loader), warmup_epochs=5)
self.evaluator = Evaluator(self.nclass, self.logger)
self.saver.save_model_summary(self.model)
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
checkpoint = torch.load(args.resume, map_location='cuda:0')
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
if args.ft:
args.start_epoch = 0
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# PATH = args.path
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args, **kwargs)
# Define network
model = SCNN(nclass=self.nclass,backbone=args.backbone,output_stride=args.out_stride,cuda = args.cuda)
# Define Optimizer
optimizer = torch.optim.SGD(model.parameters(),args.lr, momentum=args.momentum,
weight_decay=args.weight_decay, nesterov=args.nesterov)
# Define Criterion
weight = None
self.criterion = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
# patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
def __init__(self, args, ori_img_lst, init_mask_lst):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.ori_img_lst = ori_img_lst
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.test_loader, self.nclass = make_data_loader_demo(
args, args.test_folder, ori_img_lst, init_mask_lst, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn,
use_iou=args.use_maskiou)
self.model = model
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'],
strict=False)
else:
self.model.load_state_dict(checkpoint['state_dict'],
strict=False)
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def __init__(self, para):
self.args = para
# Define Saver
self.saver = Saver(para)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.train_loader, self.val_loader, self.test_loader, self.nclass = dataloader(
para)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=para.backbone,
output_stride=para.out_stride,
sync_bn=para.sync_bn,
freeze_bn=para.freeze_bn)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': para.lr
}, {
'params': model.get_10x_lr_params(),
'lr': para.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=para.momentum,
weight_decay=para.weight_decay,
nesterov=para.nesterov)
# Define Criterion
self.criterion = SegmentationLosses(
weight=None, cuda=True).build_loss(mode=para.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(para.lr_scheduler, para.lr, para.epochs,
len(self.train_loader))
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
class Trainer(object):
def __init__(self,args):
warnings.filterwarnings('ignore')
assert torch.cuda.is_available()
torch.backends.cudnn.benchmark = True
model_fname = 'data/deeplab_{0}_{1}_v3_{2}_epoch%d.pth'.format(args.backbone, args.dataset, args.exp)
if args.dataset == 'pascal':
raise NotImplementedError
elif args.dataset == 'cityscapes':
kwargs = {'num_workers': args.workers, 'pin_memory': True, 'drop_last': True}
dataset_loader, num_classes = dataloaders.make_data_loader(args, **kwargs)
args.num_classes = num_classes
elif args.dataset == 'marsh' :
kwargs = {'num_workers': args.workers, 'pin_memory': True, 'drop_last': True}
dataset_loader,val_loader, test_loader, num_classes = dataloaders.make_data_loader(args, **kwargs)
args.num_classes = num_classes
else:
raise ValueError('Unknown dataset: {}'.format(args.dataset))
if args.backbone == 'autodeeplab':
model = Retrain_Autodeeplab(args)
model.load_state_dict(torch.load(r"./run/marsh/deeplab-autodeeplab/model_best.pth.tar")['state_dict'], strict=False)
else:
raise ValueError('Unknown backbone: {}'.format(args.backbone))
optimizer = optim.SGD(model.module.parameters(), lr=args.base_lr, momentum=0.9, weight_decay=0.0001)
if args.criterion == 'Ohem':
args.thresh = 0.7
args.crop_size = [args.crop_size, args.crop_size] if isinstance(args.crop_size, int) else args.crop_size
args.n_min = int((args.batch_size / len(args.gpu) * args.crop_size[0] * args.crop_size[1]) // 16)
criterion = build_criterion(args)
model = nn.DataParallel(model).cuda()
##mergee
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
#kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = dataset_loader,val_loader, test_loader, num_classes
self.criterion = criterion
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
#self.scheduler = scheduler
self.scheduler = LR_Scheduler("poly",args.lr, args.epochs, len(self.train_loader)) #removed None from second parameter.
def __init__(self, weight_path, resume, gpu_id):
init_seeds(1)
init_dirs("result")
self.device = gpu.select_device(gpu_id)
self.start_epoch = 0
self.best_mIoU = 0.
self.epochs = cfg.TRAIN["EPOCHS"]
self.weight_path = weight_path
self.train_loader, self.val_loader, _, self.num_class = make_data_loader(
)
self.model = DeepLab(num_classes=self.num_class,
backbone="resnet",
output_stride=16,
sync_bn=False,
freeze_bn=False).to(self.device)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': cfg.TRAIN["LR_INIT"]
}, {
'params': self.model.get_10x_lr_params(),
'lr': cfg.TRAIN["LR_INIT"] * 10
}]
self.optimizer = optim.SGD(train_params,
momentum=cfg.TRAIN["MOMENTUM"],
weight_decay=cfg.TRAIN["WEIGHT_DECAY"])
self.criterion = SegmentationLosses().build_loss(
mode=cfg.TRAIN["LOSS_TYPE"])
self.scheduler = LR_Scheduler(mode=cfg.TRAIN["LR_SCHEDULER"],
base_lr=cfg.TRAIN["LR_INIT"],
num_epochs=self.epochs,
iters_per_epoch=len(self.train_loader))
self.evaluator = Evaluator(self.num_class)
self.saver = Saver()
self.summary = TensorboardSummary(os.path.join("result", "run"))
if resume:
self.__resume_model_weights()
def initialize_model(self, args):
self.args = args
# Define Saver
self.saver = Saver(self.args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_worker': self.args.worker, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
self.args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=self.args.backbone,
output_stride=self.args.out_stride,
sync_bn=self.args.sync_bn,
freeze_bn=self.args.freeze_bn)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Using cuda
if self.args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
if not os.path.isfile(self.args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
self.args.resume))
checkpoint = torch.load(self.args.resume)
self.args.start_epoch = checkpoint['epoch']
if self.args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
self.model.eval()
self.evaluator.reset()
def evaluate(self):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(self.test_loader):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
self.evaluator.add_batch(target, pred)
TensorboardSummary.visualize_images_to_folder(
self.visualizations_folder, i,
image.cpu().numpy(), target, pred, self.args.dataset)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
print('\nEvaluation:')
print('[numImages: %5d]' %
(i * self.args.batch_size + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
def testing_entropy(self):
self.saver = Saver(self.args)
self.saver.save_experiment_config()
""" Define Tensorboard Summary """
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.model.eval()
self.evaluator_1.reset()
self.evaluator_2.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output_1, avg_confidence, max_confidence = self.model.forward_testing_entropy(
image)
loss_1 = self.criterion(output_1, target)
entropy = normalized_shannon_entropy(output_1)
self.writer.add_scalar('avg_confidence/i', avg_confidence.item(),
i)
self.writer.add_scalar('max_confidence/i', max_confidence.item(),
i)
self.writer.add_scalar('entropy/i', entropy.item(), i)
self.writer.add_scalar('loss/i', loss_1.item(), i)
self.summary.visualize_image(self.writer, self.args.dataset, image,
target_show, output_2, global_step)
print('testing confidence')
self.writer.close()
def __init__(self, args, model,trn_loader,val_loader,chk_loader,optimizer):
self.args = args
self.model=model
self.train_loader = trn_loader
self.val_loader = val_loader
self.chk_loader = chk_loader
self.optimizer = optimizer
# Define Evaluator
self.evaluator = Evaluator(args.nclass)
# Define lr scheduler
# self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
# args.epochs, len(trn_loader))
self.scheduler=torch.optim.lr_scheduler.MultiStepLR(self.optimizer,milestones=[3,6,9], gamma=0.5)
self.wait_epoches=10
self.best_pred = 0
self.init_weight=0.98
# Define Saver
self.saver = Saver(self.args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.evaluator = Evaluator(self.args.nclass)
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
if not args.test:
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
if self.args.norm == 'gn':
norm = gn
elif self.args.norm == 'bn':
if self.args.sync_bn:
norm = syncbn
else:
norm = bn
elif self.args.norm == 'abn':
if self.args.sync_bn:
norm = syncabn(self.args.gpu_ids)
else:
norm = abn
else:
print("Please check the norm.")
exit()
# Define network
if self.args.model == 'deeplabv3+':
model = DeepLab(args=self.args, num_classes=self.nclass)
elif self.args.model == 'deeplabv3':
model = DeepLabv3(Norm=self.args.norm,
backbone=args.backbone,
output_stride=args.out_stride,
num_classes=self.nclass,
freeze_bn=args.freeze_bn)
elif self.args.model == 'fpn':
model = FPN(args=args, num_classes=self.nclass)
'''
model.cuda()
summary(model, input_size=(3, 720, 1280))
exit()
'''
self.classifier = Classifier(self.nclass)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset,
self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
#patch_replication_callback(self.model)
self.model = self.model.cuda()
self.classifier = torch.nn.DataParallel(
self.classifier, device_ids=self.args.gpu_ids)
self.classifier = self.classifier.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.ft:
args.start_epoch = 0
else:
args.start_epoch = checkpoint['epoch']
if args.cuda:
#self.model.module.load_state_dict(checkpoint['state_dict'])
pretrained_dict = checkpoint['state_dict']
model_dict = {}
state_dict = self.model.module.state_dict()
for k, v in pretrained_dict.items():
if k in state_dict:
model_dict[k] = v
state_dict.update(model_dict)
self.model.module.load_state_dict(state_dict)
else:
#self.model.load_state_dict(checkpoint['state_dict'])
pretrained_dict = checkpoint['state_dict']
model_dict = {}
state_dict = self.model.state_dict()
for k, v in pretrained_dict.items():
if k in state_dict:
model_dict[k] = v
state_dict.update(model_dict)
self.model.load_state_dict(state_dict)
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
elif args.decoder is not None:
if not os.path.isfile(args.decoder):
raise RuntimeError(
"=> no checkpoint for decoder found at '{}'".format(
args.decoder))
checkpoint = torch.load(args.decoder)
args.start_epoch = 0 # As every time loads decoder only should be finetuning
if args.cuda:
decoder_dict = checkpoint['state_dict']
model_dict = {}
state_dict = self.model.module.state_dict()
for k, v in decoder_dict.items():
if not 'aspp' in k:
continue
if k in state_dict:
model_dict[k] = v
state_dict.update(model_dict)
self.model.module.load_state_dict(state_dict)
else:
raise NotImplementedError("Please USE CUDA!!!")
if args.classifier is None:
raise NotImplementedError("Classifier should be loaded")
else:
if not os.path.isfile(args.classifier):
raise RuntimeError(
"=> no checkpoint for clasifier found at '{}'".format(
args.classifier))
checkpoint = torch.load(args.classifier)
s_dict = checkpoint['state_dict']
model_dict = {}
state_dict = self.classifier.state_dict()
for k, v in s_dict.items():
if k in state_dict:
model_dict[k] = v
state_dict.update(model_dict)
self.classifier.load_state_dict(state_dict)
print("Classifier checkpoint successfully loaded")
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
cell_path = os.path.join(args.saved_arch_path, 'genotype.npy')
network_path_space = os.path.join(args.saved_arch_path,
'network_path_space.npy')
new_cell_arch = np.load(cell_path)
new_network_arch = np.load(network_path_space)
# Define network
model = newModel(network_arch=new_network_arch,
cell_arch=new_cell_arch,
num_classes=self.nclass,
num_layers=12)
# output_stride=args.out_stride,
# sync_bn=args.sync_bn,
# freeze_bn=args.freeze_bn)
self.decoder = Decoder(self.nclass, 'autodeeplab', args, False)
# TODO: look into these
# TODO: ALSO look into different param groups as done int deeplab below
# train_params = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
# {'params': model.get_10x_lr_params(), 'lr': args.lr * 10}]
#
train_params = [{'params': model.parameters(), 'lr': args.lr}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset,
self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(
args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader)) #TODO: use min_lr ?
# TODO: Figure out if len(self.train_loader) should be devided by two ? in other module as well
# Using cuda
if args.cuda:
if (torch.cuda.device_count() > 1 or args.load_parallel):
self.model = torch.nn.DataParallel(self.model.cuda())
patch_replication_callback(self.model)
self.model = self.model.cuda()
print('cuda finished')
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
# if the weights are wrapped in module object we have to clean it
if args.clean_module:
self.model.load_state_dict(checkpoint['state_dict'])
state_dict = checkpoint['state_dict']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove 'module.' of dataparallel
new_state_dict[name] = v
self.model.load_state_dict(new_state_dict)
else:
if (torch.cuda.device_count() > 1 or args.load_parallel):
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
class trainNew(object):
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
cell_path = os.path.join(args.saved_arch_path, 'genotype.npy')
network_path_space = os.path.join(args.saved_arch_path,
'network_path_space.npy')
new_cell_arch = np.load(cell_path)
new_network_arch = np.load(network_path_space)
# Define network
model = newModel(network_arch=new_network_arch,
cell_arch=new_cell_arch,
num_classes=self.nclass,
num_layers=12)
# output_stride=args.out_stride,
# sync_bn=args.sync_bn,
# freeze_bn=args.freeze_bn)
self.decoder = Decoder(self.nclass, 'autodeeplab', args, False)
# TODO: look into these
# TODO: ALSO look into different param groups as done int deeplab below
# train_params = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
# {'params': model.get_10x_lr_params(), 'lr': args.lr * 10}]
#
train_params = [{'params': model.parameters(), 'lr': args.lr}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset,
self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(
args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader)) #TODO: use min_lr ?
# TODO: Figure out if len(self.train_loader) should be devided by two ? in other module as well
# Using cuda
if args.cuda:
if (torch.cuda.device_count() > 1 or args.load_parallel):
self.model = torch.nn.DataParallel(self.model.cuda())
patch_replication_callback(self.model)
self.model = self.model.cuda()
print('cuda finished')
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
# if the weights are wrapped in module object we have to clean it
if args.clean_module:
self.model.load_state_dict(checkpoint['state_dict'])
state_dict = checkpoint['state_dict']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove 'module.' of dataparallel
new_state_dict[name] = v
self.model.load_state_dict(new_state_dict)
else:
if (torch.cuda.device_count() > 1 or args.load_parallel):
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
encoder_output, low_level_feature = self.model(image)
output = self.decoder(encoder_output, low_level_feature)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(),
i + num_img_tr * epoch)
# Show 10 * 3 inference results each epoch
if i % (num_img_tr // 10) == 0:
global_step = i + num_img_tr * epoch
self.summary.visualize_image(self.writer, self.args.dataset,
image, target, output,
global_step)
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
if self.args.no_val:
# save checkpoint every epoch
is_best = False
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
encoder_output, low_level_feature = self.model(image)
output = self.decoder(encoder_output, low_level_feature)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
self.criterion = SegmentationLosses(cuda=args.cuda)
self.model, self.optimizer = model, optimizer
self.contexts = TemporalContexts(history_len=5)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning or in validation/test mode
if args.ft or args.mode == "val" or args.mode == "test":
args.start_epoch = 0
self.best_pred = 0.0
class Trainer(object):
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
print(self.nclass, args.backbone, args.out_stride, args.sync_bn,
args.freeze_bn)
#2 resnet 16 False False
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset,
self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
#image, target = sample['image'], sample['label']
image, target = sample['trace'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(),
i + num_img_tr * epoch)
# Show 10 * 3 inference results each epoch
if i % (num_img_tr // 10) == 0:
global_step = i + num_img_tr * epoch
#self.summary.visualize_image(self.writer, self.args.dataset, image, target, output, global_step)
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
if self.args.no_val:
# save checkpoint every epoch
is_best = False
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['trace'], sample['label']
#image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
class Trainer(object):
def __init__(self, args):
self.args = args
self.train_dir = './data_list/train_lite.csv'
self.train_list = pd.read_csv(self.train_dir)
self.val_dir = './data_list/val_lite.csv'
self.val_list = pd.read_csv(self.val_dir)
self.train_length = len(self.train_list)
self.val_length = len(self.val_list)
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# 方式2
self.train_gen, self.val_gen, self.test_gen, self.nclass = make_data_loader2(args)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
{'params': model.get_10x_lr_params(), 'lr': args.lr * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=args.momentum,
weight_decay=args.weight_decay, nesterov=args.nesterov)
# optimizer = torch.optim.Adam(train_params, weight_decay=args.weight_decay)
# Define Criterion
# self.criterion = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type)
self.criterion1 = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode='ce')
self.criterion2= SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode='dice')
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, self.train_length)
# Using cuda
if args.cuda:
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
# self.model.module.load_state_dict(checkpoint['state_dict'])
self.model.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def training(self, epoch):
train_loss = 0.0
prev_time = time.time()
self.model.train()
self.evaluator.reset()
num_img_tr = self.train_length / self.args.batch_size
for iteration in range(int(num_img_tr)):
samples = next(self.train_gen)
image, target = samples['image'], samples['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, iteration, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss1 = self.criterion1(output, target)
loss2 = self.criterion2(output, make_one_hot(target.long(), num_classes=self.nclass))
loss = loss1 + loss2
loss.backward()
self.optimizer.step()
train_loss += loss.item()
self.writer.add_scalar('train/total_loss_iter', loss.item(), iteration + num_img_tr * epoch)
# print log 默认log_iters = 4
if iteration % 4 == 0:
end_time = time.time()
print("Iter - %d: train loss: %.3f, celoss: %.4f, diceloss: %.4f, time cost: %.3f s" \
% (iteration, loss.item(), loss1.item(), loss2.item(), end_time - prev_time))
prev_time = time.time()
# Show 10 * 3 inference results each epoch
if iteration % (num_img_tr // 10) == 0:
global_step = iteration + num_img_tr * epoch
self.summary.visualize_image(self.writer, self.args.dataset, image, target, output, global_step)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
print("input image shape/iter:", image.shape)
# train evaluate
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
IoU = self.evaluator.Mean_Intersection_over_Union()
mIoU = np.nanmean(IoU)
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
print("Acc_tr:{}, Acc_class_tr:{}, IoU_tr:{}, mIoU_tr:{}, fwIoU_tr: {}".format(Acc, Acc_class, IoU, mIoU, FWIoU))
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' % (epoch, iteration * self.args.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
if self.args.no_val:
# save checkpoint every epoch
is_best = False
self.saver.save_checkpoint({
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
val_loss = 0.0
prev_time = time.time()
num_img_val = self.val_length / self.args.batch_size
print("Validation:","epoch ", epoch)
print(num_img_val)
for iteration in range(int(num_img_val)):
samples = next(self.val_gen)
image, target = samples['image'], samples['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad(): #
output = self.model(image)
loss1 = self.criterion1(output, target)
loss2 = self.criterion2(output, make_one_hot(target.long(), num_classes=self.nclass))
loss = loss1 + loss2
val_loss += loss.item()
self.writer.add_scalar('val/total_loss_iter', loss.item(), iteration + num_img_val * epoch)
val_loss += loss.item()
# print log 默认log_iters = 4
if iteration % 4 == 0:
end_time = time.time()
print("Iter - %d: validation loss: %.3f, celoss: %.4f, diceloss: %.4f, time cost: %.3f s" \
% (iteration, loss.item(), loss1.item(), loss2.item(), end_time - prev_time))
prev_time = time.time()
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
print(image.shape)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
IoU = self.evaluator.Mean_Intersection_over_Union()
mIoU = np.nanmean(IoU)
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', val_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' % (epoch, iteration * self.args.batch_size + image.data.shape[0]))
print("Acc_val:{}, Acc_class_val:{}, IoU:val:{}, mIoU_val:{}, fwIoU_val: {}".format(Acc, Acc_class, IoU, mIoU, FWIoU))
print('Loss: %.3f' % val_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
self.saver.save_checkpoint({
'epoch': epoch + 1,
'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def __init__(self, config, args):
self.args = args
self.config = config
self.visdom = args.visdom
if args.visdom:
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1], port=8888)
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
self.target_train_loader, self.target_val_loader, self.target_test_loader, _ = make_target_data_loader(
config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.D = Discriminator(num_classes=self.nclass, ndf=16)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': self.model.get_10x_lr_params(),
'lr': config.lr * config.lr_ratio
}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.D_optimizer = torch.optim.Adam(self.D.parameters(),
lr=config.lr,
betas=(0.9, 0.99))
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.entropy_mini_loss = MinimizeEntropyLoss()
self.bottleneck_loss = BottleneckLoss()
self.instance_loss = InstanceLoss()
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# labels for adversarial training
self.source_label = 0
self.target_label = 1
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.D = torch.nn.DataParallel(self.D)
patch_replication_callback(self.D)
self.D = self.D.cuda()
self.best_pred_source = 0.0
self.best_pred_target = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
class Trainer(object):
def __init__(self, config, args):
self.args = args
self.config = config
self.visdom = args.visdom
if args.visdom:
self.vis = visdom.Visdom(env=os.getcwd().split('/')[-1], port=8888)
# Define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
config)
self.target_train_loader, self.target_val_loader, self.target_test_loader, _ = make_target_data_loader(
config)
# Define network
self.model = DeepLab(num_classes=self.nclass,
backbone=config.backbone,
output_stride=config.out_stride,
sync_bn=config.sync_bn,
freeze_bn=config.freeze_bn)
self.D = Discriminator(num_classes=self.nclass, ndf=16)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': config.lr
}, {
'params': self.model.get_10x_lr_params(),
'lr': config.lr * config.lr_ratio
}]
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=config.momentum,
weight_decay=config.weight_decay)
self.D_optimizer = torch.optim.Adam(self.D.parameters(),
lr=config.lr,
betas=(0.9, 0.99))
# Define Criterion
# whether to use class balanced weights
self.criterion = SegmentationLosses(
weight=None, cuda=args.cuda).build_loss(mode=config.loss)
self.entropy_mini_loss = MinimizeEntropyLoss()
self.bottleneck_loss = BottleneckLoss()
self.instance_loss = InstanceLoss()
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(config.lr_scheduler,
config.lr, config.epochs,
len(self.train_loader), config.lr_step,
config.warmup_epochs)
self.summary = TensorboardSummary('./train_log')
# labels for adversarial training
self.source_label = 0
self.target_label = 1
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
# cudnn.benchmark = True
self.model = self.model.cuda()
self.D = torch.nn.DataParallel(self.D)
patch_replication_callback(self.D)
self.D = self.D.cuda()
self.best_pred_source = 0.0
self.best_pred_target = 0.0
# Resuming checkpoint
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
if args.cuda:
self.model.module.load_state_dict(checkpoint)
else:
self.model.load_state_dict(checkpoint,
map_location=torch.device('cpu'))
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, args.start_epoch))
def training(self, epoch):
train_loss, seg_loss_sum, bn_loss_sum, entropy_loss_sum, adv_loss_sum, d_loss_sum, ins_loss_sum = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
self.model.train()
if config.freeze_bn:
self.model.module.freeze_bn()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
target_train_iterator = iter(self.target_train_loader)
for i, sample in enumerate(tbar):
itr = epoch * len(self.train_loader) + i
#if self.visdom:
# self.vis.line(X=torch.tensor([itr]), Y=torch.tensor([self.optimizer.param_groups[0]['lr']]),
# win='lr', opts=dict(title='lr', xlabel='iter', ylabel='lr'),
# update='append' if itr>0 else None)
self.summary.writer.add_scalar(
'Train/lr', self.optimizer.param_groups[0]['lr'], itr)
A_image, A_target = sample['image'], sample['label']
# Get one batch from target domain
try:
target_sample = next(target_train_iterator)
except StopIteration:
target_train_iterator = iter(self.target_train_loader)
target_sample = next(target_train_iterator)
B_image, B_target, B_image_pair = target_sample[
'image'], target_sample['label'], target_sample['image_pair']
if self.args.cuda:
A_image, A_target = A_image.cuda(), A_target.cuda()
B_image, B_target, B_image_pair = B_image.cuda(
), B_target.cuda(), B_image_pair.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred_source,
self.best_pred_target, self.config.lr_ratio)
self.scheduler(self.D_optimizer, i, epoch, self.best_pred_source,
self.best_pred_target, self.config.lr_ratio)
A_output, A_feat, A_low_feat = self.model(A_image)
B_output, B_feat, B_low_feat = self.model(B_image)
#B_output_pair, B_feat_pair, B_low_feat_pair = self.model(B_image_pair)
#B_output_pair, B_feat_pair, B_low_feat_pair = flip(B_output_pair, dim=-1), flip(B_feat_pair, dim=-1), flip(B_low_feat_pair, dim=-1)
self.optimizer.zero_grad()
self.D_optimizer.zero_grad()
# Train seg network
for param in self.D.parameters():
param.requires_grad = False
# Supervised loss
seg_loss = self.criterion(A_output, A_target)
main_loss = seg_loss
# Unsupervised loss
#ins_loss = 0.01 * self.instance_loss(B_output, B_output_pair)
#main_loss += ins_loss
# Train adversarial loss
D_out = self.D(prob_2_entropy(F.softmax(B_output)))
adv_loss = bce_loss(D_out, self.source_label)
main_loss += self.config.lambda_adv * adv_loss
main_loss.backward()
# Train discriminator
for param in self.D.parameters():
param.requires_grad = True
A_output_detach = A_output.detach()
B_output_detach = B_output.detach()
# source
D_source = self.D(prob_2_entropy(F.softmax(A_output_detach)))
source_loss = bce_loss(D_source, self.source_label)
source_loss = source_loss / 2
# target
D_target = self.D(prob_2_entropy(F.softmax(B_output_detach)))
target_loss = bce_loss(D_target, self.target_label)
target_loss = target_loss / 2
d_loss = source_loss + target_loss
d_loss.backward()
self.optimizer.step()
self.D_optimizer.step()
seg_loss_sum += seg_loss.item()
#ins_loss_sum += ins_loss.item()
adv_loss_sum += self.config.lambda_adv * adv_loss.item()
d_loss_sum += d_loss.item()
#train_loss += seg_loss.item() + self.config.lambda_adv * adv_loss.item()
train_loss += seg_loss.item()
self.summary.writer.add_scalar('Train/SegLoss', seg_loss.item(),
itr)
#self.summary.writer.add_scalar('Train/InsLoss', ins_loss.item(), itr)
self.summary.writer.add_scalar('Train/AdvLoss', adv_loss.item(),
itr)
self.summary.writer.add_scalar('Train/DiscriminatorLoss',
d_loss.item(), itr)
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
# Show the results of the last iteration
#if i == len(self.train_loader)-1:
print("Add Train images at epoch" + str(epoch))
self.summary.visualize_image('Train-Source', self.config.dataset,
A_image, A_target, A_output, epoch, 5)
self.summary.visualize_image('Train-Target', self.config.target,
B_image, B_target, B_output, epoch, 5)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + A_image.data.shape[0]))
print('Loss: %.3f' % train_loss)
#print('Seg Loss: %.3f' % seg_loss_sum)
#print('Ins Loss: %.3f' % ins_loss_sum)
#print('BN Loss: %.3f' % bn_loss_sum)
#print('Adv Loss: %.3f' % adv_loss_sum)
#print('Discriminator Loss: %.3f' % d_loss_sum)
#if self.visdom:
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([seg_loss_sum]), win='train_loss', name='Seg_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([ins_loss_sum]), win='train_loss', name='Ins_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([bn_loss_sum]), win='train_loss', name='BN_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([adv_loss_sum]), win='train_loss', name='Adv_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
#self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([d_loss_sum]), win='train_loss', name='Dis_loss',
# opts=dict(title='loss', xlabel='epoch', ylabel='loss'),
# update='append' if epoch > 0 else None)
def validation(self, epoch):
def get_metrics(tbar, if_source=False):
self.evaluator.reset()
test_loss = 0.0
#feat_mean, low_feat_mean, feat_var, low_feat_var = 0, 0, 0, 0
#adv_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output, low_feat, feat = self.model(image)
#low_feat = low_feat.cpu().numpy()
#feat = feat.cpu().numpy()
#if isinstance(feat, np.ndarray):
# feat_mean += feat.mean(axis=0).mean(axis=1).mean(axis=1)
# low_feat_mean += low_feat.mean(axis=0).mean(axis=1).mean(axis=1)
# feat_var += feat.var(axis=0).var(axis=1).var(axis=1)
# low_feat_var += low_feat.var(axis=0).var(axis=1).var(axis=1)
#else:
# feat_mean = feat.mean(axis=0).mean(axis=1).mean(axis=1)
# low_feat_mean = low_feat.mean(axis=0).mean(axis=1).mean(axis=1)
# feat_var = feat.var(axis=0).var(axis=1).var(axis=1)
# low_feat_var = low_feat.var(axis=0).var(axis=1).var(axis=1)
#d_output = self.D(prob_2_entropy(F.softmax(output)))
#adv_loss += bce_loss(d_output, self.source_label).item()
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target_ = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target_, pred)
if if_source:
print("Add Validation-Source images at epoch" + str(epoch))
self.summary.visualize_image('Val-Source', self.config.dataset,
image, target, output, epoch, 5)
else:
print("Add Validation-Target images at epoch" + str(epoch))
self.summary.visualize_image('Val-Target', self.config.target,
image, target, output, epoch, 5)
#feat_mean /= (i+1)
#low_feat_mean /= (i+1)
#feat_var /= (i+1)
#low_feat_var /= (i+1)
#adv_loss /= (i+1)
# Fast test during the training
Acc = self.evaluator.Building_Acc()
IoU = self.evaluator.Building_IoU()
mIoU = self.evaluator.Mean_Intersection_over_Union()
if if_source:
print('Validation on source:')
else:
print('Validation on target:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.config.batch_size + image.data.shape[0]))
print("Acc:{}, IoU:{}, mIoU:{}".format(Acc, IoU, mIoU))
print('Loss: %.3f' % test_loss)
if if_source:
names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
self.summary.writer.add_scalar('Val/SourceAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/SourceIoU', IoU, epoch)
else:
names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
self.summary.writer.add_scalar('Val/TargetAcc', Acc, epoch)
self.summary.writer.add_scalar('Val/TargetIoU', IoU, epoch)
# Draw Visdom
#if if_source:
# names = ['source', 'source_acc', 'source_IoU', 'source_mIoU']
#else:
# names = ['target', 'target_acc', 'target_IoU', 'target_mIoU']
#if self.visdom:
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([test_loss]), win='val_loss', name=names[0],
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([adv_loss]), win='val_loss', name='adv_loss',
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([Acc]), win='metrics', name=names[1],
# opts=dict(title='metrics', xlabel='epoch', ylabel='performance'),
# update='append' if epoch > 0 else None)
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([IoU]), win='metrics', name=names[2],
# update='append')
# self.vis.line(X=torch.tensor([epoch]), Y=torch.tensor([mIoU]), win='metrics', name=names[3],
# update='append')
return Acc, IoU, mIoU
self.model.eval()
tbar_source = tqdm(self.val_loader, desc='\r')
tbar_target = tqdm(self.target_val_loader, desc='\r')
s_acc, s_iou, s_miou = get_metrics(tbar_source, True)
t_acc, t_iou, t_miou = get_metrics(tbar_target, False)
new_pred_source = s_iou
new_pred_target = t_iou
if new_pred_source > self.best_pred_source or new_pred_target > self.best_pred_target:
is_best = True
self.best_pred_source = max(new_pred_source, self.best_pred_source)
self.best_pred_target = max(new_pred_target, self.best_pred_target)
print('Saving state, epoch:', epoch)
torch.save(
self.model.module.state_dict(),
self.args.save_folder + 'models/' + 'epoch' + str(epoch) + '.pth')
loss_file = {
's_Acc': s_acc,
's_IoU': s_iou,
's_mIoU': s_miou,
't_Acc': t_acc,
't_IoU': t_iou,
't_mIoU': t_miou
}
with open(
os.path.join(self.args.save_folder, 'eval',
'epoch' + str(epoch) + '.json'), 'w') as f:
json.dump(loss_file, f)
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
{'params': model.get_10x_lr_params(), 'lr': args.lr * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=args.momentum,
weight_decay=args.weight_decay, nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(Path.db_root_dir(args.dataset), args.dataset+'_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
class Trainer(object):
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
{'params': model.get_10x_lr_params(), 'lr': args.lr * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=args.momentum,
weight_decay=args.weight_decay, nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(Path.db_root_dir(args.dataset), args.dataset+'_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model, device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'" .format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch)
# Show 10 * 3 inference results each epoch
if i % (num_img_tr // 10) == 0:
global_step = i + num_img_tr * epoch
self.summary.visualize_image(self.writer, self.args.dataset, image, target, output, global_step)
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
if self.args.no_val:
# save checkpoint every epoch
is_best = False
self.saver.save_checkpoint({
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
self.saver.save_checkpoint({
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
class Trainer(object):
def __init__(self, config):
self.config = config
self.best_pred = 0.0
# Define Saver
self.saver = Saver(config)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.config['training']['tensorboard']['log_dir'])
self.writer = self.summary.create_summary()
self.train_loader, self.val_loader, self.test_loader, self.nclass = initialize_data_loader(config)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=self.config['network']['backbone'],
output_stride=self.config['image']['out_stride'],
sync_bn=self.config['network']['sync_bn'],
freeze_bn=self.config['network']['freeze_bn'])
train_params = [{'params': model.get_1x_lr_params(), 'lr': self.config['training']['lr']},
{'params': model.get_10x_lr_params(), 'lr': self.config['training']['lr'] * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=self.config['training']['momentum'],
weight_decay=self.config['training']['weight_decay'], nesterov=self.config['training']['nesterov'])
# Define Criterion
# whether to use class balanced weights
if self.config['training']['use_balanced_weights']:
classes_weights_path = os.path.join(self.config['dataset']['base_path'], self.config['dataset']['dataset_name'] + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(self.config, self.config['dataset']['dataset_name'], self.train_loader, self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(weight=weight, cuda=self.config['network']['use_cuda']).build_loss(mode=self.config['training']['loss_type'])
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(self.config['training']['lr_scheduler'], self.config['training']['lr'],
self.config['training']['epochs'], len(self.train_loader))
# Using cuda
if self.config['network']['use_cuda']:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
if self.config['training']['weights_initialization']['use_pretrained_weights']:
if not os.path.isfile(self.config['training']['weights_initialization']['restore_from']):
raise RuntimeError("=> no checkpoint found at '{}'" .format(self.config['training']['weights_initialization']['restore_from']))
if self.config['network']['use_cuda']:
checkpoint = torch.load(self.config['training']['weights_initialization']['restore_from'])
else:
checkpoint = torch.load(self.config['training']['weights_initialization']['restore_from'], map_location={'cuda:0': 'cpu'})
self.config['training']['start_epoch'] = checkpoint['epoch']
if self.config['network']['use_cuda']:
self.model.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
# if not self.config['ft']:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(self.config['training']['weights_initialization']['restore_from'], checkpoint['epoch']))
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.config['network']['use_cuda']:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch)
# Show 10 * 3 inference results each epoch
if i % (num_img_tr // 10) == 0:
global_step = i + num_img_tr * epoch
self.summary.visualize_image(self.writer, self.config['dataset']['dataset_name'], image, target, output, global_step)
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.config['training']['batch_size'] + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
#save last checkpoint
self.saver.save_checkpoint({
'epoch': epoch + 1,
# 'state_dict': self.model.module.state_dict(),
'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best = False, filename='checkpoint_last.pth.tar')
#if training on a subset reshuffle the data
if self.config['training']['train_on_subset']['enabled']:
self.train_loader.dataset.shuffle_dataset()
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.config['network']['use_cuda']:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Val loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.config['training']['batch_size'] + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
self.best_pred = new_pred
self.saver.save_checkpoint({
'epoch': epoch + 1,
# 'state_dict': self.model.module.state_dict(),
'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best = True, filename='checkpoint_best.pth.tar')
def __init__(self, config):
self.config = config
self.best_pred = 0.0
# Define Saver
self.saver = Saver(config)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.config['training']['tensorboard']['log_dir'])
self.writer = self.summary.create_summary()
self.train_loader, self.val_loader, self.test_loader, self.nclass = initialize_data_loader(config)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=self.config['network']['backbone'],
output_stride=self.config['image']['out_stride'],
sync_bn=self.config['network']['sync_bn'],
freeze_bn=self.config['network']['freeze_bn'])
train_params = [{'params': model.get_1x_lr_params(), 'lr': self.config['training']['lr']},
{'params': model.get_10x_lr_params(), 'lr': self.config['training']['lr'] * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=self.config['training']['momentum'],
weight_decay=self.config['training']['weight_decay'], nesterov=self.config['training']['nesterov'])
# Define Criterion
# whether to use class balanced weights
if self.config['training']['use_balanced_weights']:
classes_weights_path = os.path.join(self.config['dataset']['base_path'], self.config['dataset']['dataset_name'] + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(self.config, self.config['dataset']['dataset_name'], self.train_loader, self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(weight=weight, cuda=self.config['network']['use_cuda']).build_loss(mode=self.config['training']['loss_type'])
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(self.config['training']['lr_scheduler'], self.config['training']['lr'],
self.config['training']['epochs'], len(self.train_loader))
# Using cuda
if self.config['network']['use_cuda']:
self.model = torch.nn.DataParallel(self.model)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
if self.config['training']['weights_initialization']['use_pretrained_weights']:
if not os.path.isfile(self.config['training']['weights_initialization']['restore_from']):
raise RuntimeError("=> no checkpoint found at '{}'" .format(self.config['training']['weights_initialization']['restore_from']))
if self.config['network']['use_cuda']:
checkpoint = torch.load(self.config['training']['weights_initialization']['restore_from'])
else:
checkpoint = torch.load(self.config['training']['weights_initialization']['restore_from'], map_location={'cuda:0': 'cpu'})
self.config['training']['start_epoch'] = checkpoint['epoch']
if self.config['network']['use_cuda']:
self.model.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
# if not self.config['ft']:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(self.config['training']['weights_initialization']['restore_from'], checkpoint['epoch']))
class Trainer(object):
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.use_amp = True if (APEX_AVAILABLE and args.use_amp) else False
self.opt_level = args.opt_level
kwargs = {
'num_workers': args.workers,
'pin_memory': True,
'drop_last': True
}
self.train_loaderA, self.train_loaderB, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
raise NotImplementedError
#if so, which trainloader to use?
# weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
# Define network
model = AutoDeeplab(self.nclass, 12, self.criterion,
self.args.filter_multiplier,
self.args.block_multiplier, self.args.step)
optimizer = torch.optim.SGD(model.weight_parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.model, self.optimizer = model, optimizer
self.architect_optimizer = torch.optim.Adam(
self.model.arch_parameters(),
lr=args.arch_lr,
betas=(0.9, 0.999),
weight_decay=args.arch_weight_decay)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.train_loaderA),
min_lr=args.min_lr)
# TODO: Figure out if len(self.train_loader) should be devided by two ? in other module as well
# Using cuda
if args.cuda:
self.model = self.model.cuda()
# mixed precision
if self.use_amp and args.cuda:
keep_batchnorm_fp32 = True if (self.opt_level == 'O2'
or self.opt_level == 'O3') else None
# fix for current pytorch version with opt_level 'O1'
if self.opt_level == 'O1' and torch.__version__ < '1.3':
for module in self.model.modules():
if isinstance(module,
torch.nn.modules.batchnorm._BatchNorm):
# Hack to fix BN fprop without affine transformation
if module.weight is None:
module.weight = torch.nn.Parameter(
torch.ones(module.running_var.shape,
dtype=module.running_var.dtype,
device=module.running_var.device),
requires_grad=False)
if module.bias is None:
module.bias = torch.nn.Parameter(
torch.zeros(module.running_var.shape,
dtype=module.running_var.dtype,
device=module.running_var.device),
requires_grad=False)
# print(keep_batchnorm_fp32)
self.model, [self.optimizer,
self.architect_optimizer] = amp.initialize(
self.model,
[self.optimizer, self.architect_optimizer],
opt_level=self.opt_level,
keep_batchnorm_fp32=keep_batchnorm_fp32,
loss_scale="dynamic")
print('cuda finished')
# Using data parallel
if args.cuda and len(self.args.gpu_ids) > 1:
if self.opt_level == 'O2' or self.opt_level == 'O3':
print(
'currently cannot run with nn.DataParallel and optimization level',
self.opt_level)
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
print('training on multiple-GPUs')
#checkpoint = torch.load(args.resume)
#print('about to load state_dict')
#self.model.load_state_dict(checkpoint['state_dict'])
#print('model loaded')
#sys.exit()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
# if the weights are wrapped in module object we have to clean it
if args.clean_module:
self.model.load_state_dict(checkpoint['state_dict'])
state_dict = checkpoint['state_dict']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove 'module.' of dataparallel
new_state_dict[name] = v
# self.model.load_state_dict(new_state_dict)
copy_state_dict(self.model.state_dict(), new_state_dict)
else:
if torch.cuda.device_count() > 1 or args.load_parallel:
# self.model.module.load_state_dict(checkpoint['state_dict'])
copy_state_dict(self.model.module.state_dict(),
checkpoint['state_dict'])
else:
# self.model.load_state_dict(checkpoint['state_dict'])
copy_state_dict(self.model.state_dict(),
checkpoint['state_dict'])
if not args.ft:
# self.optimizer.load_state_dict(checkpoint['optimizer'])
copy_state_dict(self.optimizer.state_dict(),
checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loaderA)
num_img_tr = len(self.train_loaderA)
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
if self.use_amp:
with amp.scale_loss(loss, self.optimizer) as scaled_loss:
scaled_loss.backward()
else:
loss.backward()
self.optimizer.step()
if epoch >= self.args.alpha_epoch:
search = next(iter(self.train_loaderB))
image_search, target_search = search['image'], search['label']
if self.args.cuda:
image_search, target_search = image_search.cuda(
), target_search.cuda()
self.architect_optimizer.zero_grad()
output_search = self.model(image_search)
arch_loss = self.criterion(output_search, target_search)
if self.use_amp:
with amp.scale_loss(
arch_loss,
self.architect_optimizer) as arch_scaled_loss:
arch_scaled_loss.backward()
else:
arch_loss.backward()
self.architect_optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
#self.writer.add_scalar('train/total_loss_iter', loss.item(), i + num_img_tr * epoch)
# Show 10 * 3 inference results each epoch
if i % (num_img_tr // 10) == 0:
global_step = i + num_img_tr * epoch
self.summary.visualize_image(self.writer, self.args.dataset,
image, target, output,
global_step)
#torch.cuda.empty_cache()
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print('Loss: %.3f' % train_loss)
if self.args.no_val:
# save checkpoint every epoch
is_best = False
if torch.cuda.device_count() > 1:
state_dict = self.model.module.state_dict()
else:
state_dict = self.model.state_dict()
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': state_dict,
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
tbar = tqdm(self.val_loader, desc='\r')
test_loss = 0.0
for i, sample in enumerate(tbar):
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
test_loss += loss.item()
tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print('[Epoch: %d, numImages: %5d]' %
(epoch, i * self.args.batch_size + image.data.shape[0]))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
print('Loss: %.3f' % test_loss)
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
if torch.cuda.device_count() > 1:
state_dict = self.model.module.state_dict()
else:
state_dict = self.model.state_dict()
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': state_dict,
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
def __init__(self, args):
self.args = args
self.train_dir = './data_list/train_lite.csv'
self.train_list = pd.read_csv(self.train_dir)
self.val_dir = './data_list/val_lite.csv'
self.val_list = pd.read_csv(self.val_dir)
self.train_length = len(self.train_list)
self.val_length = len(self.val_list)
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# 方式2
self.train_gen, self.val_gen, self.test_gen, self.nclass = make_data_loader2(args)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{'params': model.get_1x_lr_params(), 'lr': args.lr},
{'params': model.get_10x_lr_params(), 'lr': args.lr * 10}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params, momentum=args.momentum,
weight_decay=args.weight_decay, nesterov=args.nesterov)
# optimizer = torch.optim.Adam(train_params, weight_decay=args.weight_decay)
# Define Criterion
# self.criterion = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type)
self.criterion1 = SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode='ce')
self.criterion2= SegmentationLosses(weight=None, cuda=args.cuda).build_loss(mode='dice')
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr,
args.epochs, self.train_length)
# Using cuda
if args.cuda:
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
# self.model.module.load_state_dict(checkpoint['state_dict'])
self.model.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
self.use_amp = True if (APEX_AVAILABLE and args.use_amp) else False
self.opt_level = args.opt_level
kwargs = {
'num_workers': args.workers,
'pin_memory': True,
'drop_last': True
}
self.train_loaderA, self.train_loaderB, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
raise NotImplementedError
#if so, which trainloader to use?
# weight = calculate_weigths_labels(args.dataset, self.train_loader, self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
# Define network
model = AutoDeeplab(self.nclass, 12, self.criterion,
self.args.filter_multiplier,
self.args.block_multiplier, self.args.step)
optimizer = torch.optim.SGD(model.weight_parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.model, self.optimizer = model, optimizer
self.architect_optimizer = torch.optim.Adam(
self.model.arch_parameters(),
lr=args.arch_lr,
betas=(0.9, 0.999),
weight_decay=args.arch_weight_decay)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.train_loaderA),
min_lr=args.min_lr)
# TODO: Figure out if len(self.train_loader) should be devided by two ? in other module as well
# Using cuda
if args.cuda:
self.model = self.model.cuda()
# mixed precision
if self.use_amp and args.cuda:
keep_batchnorm_fp32 = True if (self.opt_level == 'O2'
or self.opt_level == 'O3') else None
# fix for current pytorch version with opt_level 'O1'
if self.opt_level == 'O1' and torch.__version__ < '1.3':
for module in self.model.modules():
if isinstance(module,
torch.nn.modules.batchnorm._BatchNorm):
# Hack to fix BN fprop without affine transformation
if module.weight is None:
module.weight = torch.nn.Parameter(
torch.ones(module.running_var.shape,
dtype=module.running_var.dtype,
device=module.running_var.device),
requires_grad=False)
if module.bias is None:
module.bias = torch.nn.Parameter(
torch.zeros(module.running_var.shape,
dtype=module.running_var.dtype,
device=module.running_var.device),
requires_grad=False)
# print(keep_batchnorm_fp32)
self.model, [self.optimizer,
self.architect_optimizer] = amp.initialize(
self.model,
[self.optimizer, self.architect_optimizer],
opt_level=self.opt_level,
keep_batchnorm_fp32=keep_batchnorm_fp32,
loss_scale="dynamic")
print('cuda finished')
# Using data parallel
if args.cuda and len(self.args.gpu_ids) > 1:
if self.opt_level == 'O2' or self.opt_level == 'O3':
print(
'currently cannot run with nn.DataParallel and optimization level',
self.opt_level)
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
print('training on multiple-GPUs')
#checkpoint = torch.load(args.resume)
#print('about to load state_dict')
#self.model.load_state_dict(checkpoint['state_dict'])
#print('model loaded')
#sys.exit()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
# if the weights are wrapped in module object we have to clean it
if args.clean_module:
self.model.load_state_dict(checkpoint['state_dict'])
state_dict = checkpoint['state_dict']
new_state_dict = OrderedDict()
for k, v in state_dict.items():
name = k[7:] # remove 'module.' of dataparallel
new_state_dict[name] = v
# self.model.load_state_dict(new_state_dict)
copy_state_dict(self.model.state_dict(), new_state_dict)
else:
if torch.cuda.device_count() > 1 or args.load_parallel:
# self.model.module.load_state_dict(checkpoint['state_dict'])
copy_state_dict(self.model.module.state_dict(),
checkpoint['state_dict'])
else:
# self.model.load_state_dict(checkpoint['state_dict'])
copy_state_dict(self.model.state_dict(),
checkpoint['state_dict'])
if not args.ft:
# self.optimizer.load_state_dict(checkpoint['optimizer'])
copy_state_dict(self.optimizer.state_dict(),
checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
print(self.nclass, args.backbone, args.out_stride, args.sync_bn,
args.freeze_bn)
#2 resnet 16 False False
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
# whether to use class balanced weights
if args.use_balanced_weights:
classes_weights_path = os.path.join(
Path.db_root_dir(args.dataset),
args.dataset + '_classes_weights.npy')
if os.path.isfile(classes_weights_path):
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(args.dataset,
self.train_loader,
self.nclass)
weight = torch.from_numpy(weight.astype(np.float32))
else:
weight = None
self.criterion = SegmentationLosses(
weight=weight, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume, map_location='cpu')
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning
if args.ft:
args.start_epoch = 0
class Trainer(object):
def __init__(self, args):
self.args = args
# Define Saver
self.saver = Saver(args)
self.saver.save_experiment_config()
# Define Tensorboard Summary
self.summary = TensorboardSummary(self.saver.experiment_dir)
self.writer = self.summary.create_summary()
# Define Dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True}
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(
args, **kwargs)
# Define network
model = DeepLab(num_classes=self.nclass,
backbone=args.backbone,
output_stride=args.out_stride,
sync_bn=args.sync_bn,
freeze_bn=args.freeze_bn)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': args.lr
}, {
'params': model.get_10x_lr_params(),
'lr': args.lr * 10
}]
# Define Optimizer
optimizer = torch.optim.SGD(train_params,
momentum=args.momentum,
weight_decay=args.weight_decay,
nesterov=args.nesterov)
# Define Criterion
self.criterion = SegmentationLosses(cuda=args.cuda)
self.model, self.optimizer = model, optimizer
self.contexts = TemporalContexts(history_len=5)
# Define Evaluator
self.evaluator = Evaluator(self.nclass)
# Define lr scheduler
self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, args.epochs,
len(self.train_loader))
# Using cuda
if args.cuda:
self.model = torch.nn.DataParallel(self.model,
device_ids=self.args.gpu_ids)
patch_replication_callback(self.model)
self.model = self.model.cuda()
# Resuming checkpoint
self.best_pred = 0.0
if args.resume is not None:
if not os.path.isfile(args.resume):
raise RuntimeError("=> no checkpoint found at '{}'".format(
args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
if args.cuda:
self.model.module.load_state_dict(checkpoint['state_dict'])
else:
self.model.load_state_dict(checkpoint['state_dict'])
if not args.ft:
self.optimizer.load_state_dict(checkpoint['optimizer'])
self.best_pred = checkpoint['best_pred']
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# Clear start epoch if fine-tuning or in validation/test mode
if args.ft or args.mode == "val" or args.mode == "test":
args.start_epoch = 0
self.best_pred = 0.0
def training(self, epoch):
train_loss = 0.0
self.model.train()
tbar = tqdm(self.train_loader)
num_img_tr = len(self.train_loader)
for i, sample in enumerate(tbar):
image, region_prop, target = sample['image'], sample['rp'], sample[
'label']
if self.args.cuda:
image, region_prop, target = image.cuda(), region_prop.cuda(
), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image, region_prop)
loss = self.criterion.CrossEntropyLoss(
output,
target,
weight=torch.from_numpy(
calculate_weights_batch(sample,
self.nclass).astype(np.float32)))
loss.backward()
self.optimizer.step()
train_loss += loss.item()
tbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(),
i + num_img_tr * epoch)
pred = output.clone().data.cpu()
pred_softmax = F.softmax(pred, dim=1).numpy()
pred = np.argmax(pred.numpy(), axis=1)
# Plot prediction every 20th iter
if i % (num_img_tr // 20) == 0:
global_step = i + num_img_tr * epoch
self.summary.vis_grid(self.writer,
self.args.dataset,
image.data.cpu().numpy()[0],
target.data.cpu().numpy()[0],
pred[0],
region_prop.data.cpu().numpy()[0],
pred_softmax[0],
global_step,
split="Train")
self.writer.add_scalar('train/total_loss_epoch',
train_loss / num_img_tr, epoch)
print('Loss: {}'.format(train_loss / num_img_tr))
if self.args.no_val or self.args.save_all:
# save checkpoint every epoch
is_best = False
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
},
is_best,
filename='checkpoint_' + str(epoch + 1) + '_.pth.tar')
def validation(self, epoch):
if self.args.mode == "train" or self.args.mode == "val":
loader = self.val_loader
elif self.args.mode == "test":
loader = self.test_loader
self.model.eval()
self.evaluator.reset()
tbar = tqdm(loader, desc='\r')
test_loss = 0.0
idr_thresholds = [0.20, 0.30, 0.40, 0.50, 0.60, 0.65]
num_itr = len(loader)
for i, sample in enumerate(tbar):
image, region_prop, target = sample['image'], sample['rp'], sample[
'label']
# orig_region_prop = region_prop.clone()
# region_prop = self.contexts.temporal_prop(image.numpy(),region_prop.numpy())
if self.args.cuda:
image, region_prop, target = image.cuda(), region_prop.cuda(
), target.cuda()
with torch.no_grad():
output = self.model(image, region_prop)
# loss = self.criterion.CrossEntropyLoss(output,target,weight=torch.from_numpy(calculate_weights_batch(sample,self.nclass).astype(np.float32)))
# test_loss += loss.item()
# tbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
output = output.detach().data.cpu()
pred_softmax = F.softmax(output, dim=1).numpy()
pred = np.argmax(pred_softmax, axis=1)
target = target.cpu().numpy()
image = image.cpu().numpy()
region_prop = region_prop.cpu().numpy()
# orig_region_prop = orig_region_prop.numpy()
# Add batch sample into evaluator
self.evaluator.add_batch(target, pred)
# Append buffer with original context(before temporal propagation)
# self.contexts.append_buffer(image[0],orig_region_prop[0],pred[0])
global_step = i + num_itr * epoch
self.summary.vis_grid(self.writer,
self.args.dataset,
image[0],
target[0],
pred[0],
region_prop[0],
pred_softmax[0],
global_step,
split="Validation")
# Fast test during the training
mIoU = self.evaluator.Mean_Intersection_over_Union()
recall, precision = self.evaluator.pdr_metric(class_id=2)
idr_avg = np.array([
self.evaluator.get_idr(class_value=2, threshold=value)
for value in idr_thresholds
])
false_idr = self.evaluator.get_false_idr(class_value=2)
instance_iou = self.evaluator.get_instance_iou(threshold=0.20,
class_value=2)
# self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Recall/per_epoch', recall, epoch)
self.writer.add_scalar('IDR/per_epoch(0.20)', idr_avg[0], epoch)
self.writer.add_scalar('IDR/avg_epoch', np.mean(idr_avg), epoch)
self.writer.add_scalar('False_IDR/epoch', false_idr, epoch)
self.writer.add_scalar('Instance_IOU/epoch', instance_iou, epoch)
self.writer.add_histogram(
'Prediction_hist',
self.evaluator.pred_labels[self.evaluator.gt_labels == 2], epoch)
print('Validation:')
# print('Loss: %.3f' % test_loss)
# print('Recall/PDR:{}'.format(recall))
print('IDR:{}'.format(idr_avg[0]))
print('False Positive Rate: {}'.format(false_idr))
print('Instance_IOU: {}'.format(instance_iou))
if self.args.mode == "train":
new_pred = mIoU
if new_pred > self.best_pred:
is_best = True
self.best_pred = new_pred
self.saver.save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': self.model.module.state_dict(),
'optimizer': self.optimizer.state_dict(),
'best_pred': self.best_pred,
}, is_best)
else:
pass
