def __init__(self, args):
self.args = args
args.log_name = str(args.checkname)
self.logger = utils.create_logger(args.log_root, args.log_name)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'logger': self.logger,
'scale': args.scale
}
trainset = get_segmentation_dataset(args.dataset,
split='train',
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_segmentation_model(args.model,
dataset=args.dataset,
backbone=args.backbone,
aux=args.aux,
se_loss=args.se_loss,
norm_layer=BatchNorm2d,
base_size=args.base_size,
crop_size=args.crop_size,
multi_grid=args.multi_grid,
multi_dilation=args.multi_dilation)
#print(model)
self.logger.info(model)
# optimizer using different LR
params_list = [
{
'params': model.pretrained.parameters(),
'lr': args.lr
},
]
if hasattr(model, 'head'):
params_list.append({
'params': model.head.parameters(),
'lr': args.lr * 10
})
if hasattr(model, 'auxlayer'):
params_list.append({
'params': model.auxlayer.parameters(),
'lr': args.lr * 10
})
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.criterion = SegmentationMultiLosses(nclass=self.nclass)
#self.criterion = SegmentationLosses(se_loss=args.se_loss, aux=args.aux,nclass=self.nclass)
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# finetune from a trained model
if args.ft:
args.start_epoch = 0
checkpoint = torch.load(args.ft_resume)
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.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.ft_resume, checkpoint['epoch']))
# resuming checkpoint
if args.resume:
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']
self.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader),
logger=self.logger,
lr_step=args.lr_step)
self.best_pred = 0.0
def __init__(self, args):
self.args = args
self.logger = utils.create_logger(self.args.exp_dir, "log")
for k, v in vars(self.args).items():
self.logger.info((k, v))
if self.args.cuda:
device = torch.device("cuda")
self.logger.info("training on gpu:" + self.args.gpu_id)
else:
self.logger.info("training on cpu")
device = torch.device("cpu")
self.device = device
#指定随机数
set_seed(args.random_seed)
# args.log_name = str(args.checkname)
#好像是可以加速
cudnn.benchmark = True
#读取数据集，现在只有VOC
self.logger.info('training on dataset ' + self.args.dataset)
input_transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize([.485, .456, .406], [.229, .224, .225])
])
data_kwargs = {
'transform': input_transform,
'root': self.args.data_dir,
'base_size': self.args.base_size,
'crop_size': self.args.crop_size,
'logger': self.logger,
'scale': self.args.scale
}
trainset = get_segmentation_dataset(self.args.dataset,
split='train',
mode='train',
**data_kwargs)
testset = get_segmentation_dataset(self.args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': self.args.num_workers, 'pin_memory': True} \
if self.args.cuda else {}
self.train_iter = data.DataLoader(trainset,
batch_size=self.args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.val_iter = data.DataLoader(testset,
batch_size=self.args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.num_classes = trainset.num_classes
self.input_channels = trainset.input_channels
#create model
kwargs = {'fuse_attention': self.args.fuse_attention}
self.model = get_segmentation_model(args.arch,
dataset=args.dataset,
backbone=args.backbone)
print("=> creating model %s" % self.args.arch)
# self.model = archs.__dict__[self.args.arch](num_classes=self.num_classes,
# input_channels=self.input_channels, **model_kwargs)
self.model = self.model.to(device)
# self.logger.info(self.model)
self.optimizer = None
params = filter(lambda p: p.requires_grad, self.model.parameters())
if self.args.optimizer == "Adam":
self.optimizer = torch.optim.Adam(
params, lr=self.args.lr, weight_decay=self.args.weight_decay)
elif self.args.optimizer == 'SGD':
self.optimizer = torch.optim.SGD(
params,
lr=self.args.lr,
momentum=self.args.momentum,
weight_decay=self.args.weight_decay)
else:
raise NotImplementedError
#loss函数
self.criterion = nn.CrossEntropyLoss(
ignore_index=trainset.IGNORE_INDEX)
#语义分割评价指标
self.metric = SegmentationMetric(self.num_classes)
#学习率策略
self.scheduler = LR_Scheduler(self.args.scheduler, base_lr = self.args.lr, num_epochs=self.args.epochs, \
iters_per_epoch=len(self.train_iter))
#创建实验结果保存目录
self.writer = SummaryWriter(args.exp_dir)
# with open(os.path.join(args.exp_dir,'config.yml'), 'w') as f:
# yaml.dump(config, f)
#用tensoboard看一下模型结构
X, label = next(iter(self.train_iter))
self.writer.add_graph(self.model, X.to(device))
self.epoch_begin = 0
self.best_iou = 0.0
#在训练开始前看看输出是什么
val_log = self.validate(epoch=-1, is_visualize_segmentation=True)
self.write_into_tensorboard(val_log, val_log, epoch=-1)
#checkpoint_PATH
if self.args.checkpoint_PATH is not None:
if self.args.only_read_model:
model, _, _, _, _ = load_checkpoint(model,
self.args.checkpoint_PATH)
else:
model, self.epoch_begin, self.best_iou, self.optimizer = load_checkpoint(
model, self.args.checkpoint_PATH, epoch_begin, best_iou,
optimizer, scheduler)
def main(args):
writer = SummaryWriter(log_dir=args.tensorboard_log_dir)
w, h = map(int, args.input_size.split(','))
w_target, h_target = map(int, args.input_size_target.split(','))
joint_transform = joint_transforms.Compose([
joint_transforms.FreeScale((h, w)),
joint_transforms.RandomHorizontallyFlip(),
])
normalize = ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
input_transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*normalize),
])
target_transform = extended_transforms.MaskToTensor()
restore_transform = standard_transforms.ToPILImage()
if '5' in args.data_dir:
dataset = GTA5DataSetLMDB(
args.data_dir, args.data_list,
joint_transform=joint_transform,
transform=input_transform, target_transform=target_transform,
)
else:
dataset = CityscapesDataSetLMDB(
args.data_dir, args.data_list,
joint_transform=joint_transform,
transform=input_transform, target_transform=target_transform,
)
loader = data.DataLoader(
dataset, batch_size=args.batch_size,
shuffle=True, num_workers=args.num_workers, pin_memory=True
)
val_dataset = CityscapesDataSetLMDB(
args.data_dir_target, args.data_list_target,
# joint_transform=joint_transform,
transform=input_transform, target_transform=target_transform
)
val_loader = data.DataLoader(
val_dataset, batch_size=args.batch_size,
shuffle=False, num_workers=args.num_workers, pin_memory=True
)
upsample = nn.Upsample(size=(h_target, w_target),
mode='bilinear', align_corners=True)
net = PSP(
nclass = args.n_classes, backbone='resnet101',
root=args.model_path_prefix, norm_layer=BatchNorm2d,
)
params_list = [
{'params': net.pretrained.parameters(), 'lr': args.learning_rate},
{'params': net.head.parameters(), 'lr': args.learning_rate*10},
{'params': net.auxlayer.parameters(), 'lr': args.learning_rate*10},
]
optimizer = torch.optim.SGD(params_list,
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay)
criterion = SegmentationLosses(nclass=args.n_classes, aux=True, ignore_index=255)
# criterion = SegmentationMultiLosses(nclass=args.n_classes, ignore_index=255)
net = DataParallelModel(net).cuda()
criterion = DataParallelCriterion(criterion).cuda()
logger = utils.create_logger(args.tensorboard_log_dir, 'PSP_train')
scheduler = utils.LR_Scheduler(args.lr_scheduler, args.learning_rate,
args.num_epoch, len(loader), logger=logger,
lr_step=args.lr_step)
net_eval = Eval(net)
num_batches = len(loader)
best_pred = 0.0
for epoch in range(args.num_epoch):
loss_rec = AverageMeter()
data_time_rec = AverageMeter()
batch_time_rec = AverageMeter()
tem_time = time.time()
for batch_index, batch_data in enumerate(loader):
scheduler(optimizer, batch_index, epoch, best_pred)
show_fig = (batch_index+1) % args.show_img_freq == 0
iteration = batch_index+1+epoch*num_batches
net.train()
img, label, name = batch_data
img = img.cuda()
label_cuda = label.cuda()
data_time_rec.update(time.time()-tem_time)
output = net(img)
loss = criterion(output, label_cuda)
optimizer.zero_grad()
loss.backward()
optimizer.step()
loss_rec.update(loss.item())
writer.add_scalar('A_seg_loss', loss.item(), iteration)
batch_time_rec.update(time.time()-tem_time)
tem_time = time.time()
if (batch_index+1) % args.print_freq == 0:
print(
f'Epoch [{epoch+1:d}/{args.num_epoch:d}][{batch_index+1:d}/{num_batches:d}]\t'
f'Time: {batch_time_rec.avg:.2f} '
f'Data: {data_time_rec.avg:.2f} '
f'Loss: {loss_rec.avg:.2f}'
)
# if show_fig:
# # base_lr = optimizer.param_groups[0]["lr"]
# output = torch.argmax(output[0][0], dim=1).detach()[0, ...].cpu()
# # fig, axes = plt.subplots(2, 1, figsize=(12, 14))
# # axes = axes.flat
# # axes[0].imshow(colorize_mask(output.numpy()))
# # axes[0].set_title(name[0])
# # axes[1].imshow(colorize_mask(label[0, ...].numpy()))
# # axes[1].set_title(f'seg_true_{base_lr:.6f}')
# # writer.add_figure('A_seg', fig, iteration)
# output_mask = np.asarray(colorize_mask(output.numpy()))
# label = np.asarray(colorize_mask(label[0,...].numpy()))
# image_out = np.concatenate([output_mask, label])
# writer.add_image('A_seg', image_out, iteration)
mean_iu = test_miou(net_eval, val_loader, upsample,
'./style_seg/dataset/info.json')
torch.save(
net.module.state_dict(),
os.path.join(args.save_path_prefix, f'{epoch:d}_{mean_iu*100:.0f}.pth')
)
writer.close()
def __init__(self, args):
self.args = args
args.log_name = str(args.checkname)
self.logger = utils.create_logger(args.log_root, args.log_name)
self.logger.info(args)
# data transforms
input_transform = transform.Compose([
transform.ToTensor(),
transform.Normalize([.485, .456, .406], [.229, .224, .225])
])
# dataset
data_kwargs = {
'transform': input_transform,
'base_size': args.base_size,
'crop_size': args.crop_size,
'logger': self.logger,
'scale': args.scale
}
trainset = get_edge_dataset(args.dataset,
split='train',
mode='train',
**data_kwargs)
testset = get_edge_dataset(args.dataset,
split='val',
mode='val',
**data_kwargs)
# dataloader
kwargs = {'num_workers': args.workers, 'pin_memory': True} \
if args.cuda else {}
self.trainloader = data.DataLoader(trainset,
batch_size=args.batch_size,
drop_last=True,
shuffle=True,
**kwargs)
self.valloader = data.DataLoader(testset,
batch_size=args.batch_size,
drop_last=False,
shuffle=False,
**kwargs)
self.nclass = trainset.num_class
# model
model = get_edge_model(
args.model,
dataset=args.dataset,
backbone=args.backbone,
norm_layer=BatchNorm2d,
crop_size=args.crop_size,
)
self.logger.info(model)
# optimizer using different LR
if args.model == 'dff': # dff
params_list = [{
'params': model.pretrained.parameters(),
'lr': args.lr
}, {
'params': model.ada_learner.parameters(),
'lr': args.lr * 10
}, {
'params': model.side1.parameters(),
'lr': args.lr * 10
}, {
'params': model.side2.parameters(),
'lr': args.lr * 10
}, {
'params': model.side3.parameters(),
'lr': args.lr * 10
}, {
'params': model.side5.parameters(),
'lr': args.lr * 10
}, {
'params': model.side5_w.parameters(),
'lr': args.lr * 10
}]
else: # casenet
assert args.model == 'casenet'
params_list = [{
'params': model.pretrained.parameters(),
'lr': args.lr
}, {
'params': model.side1.parameters(),
'lr': args.lr * 10
}, {
'params': model.side2.parameters(),
'lr': args.lr * 10
}, {
'params': model.side3.parameters(),
'lr': args.lr * 10
}, {
'params': model.side5.parameters(),
'lr': args.lr * 10
}, {
'params': model.fuse.parameters(),
'lr': args.lr * 10
}]
optimizer = torch.optim.SGD(params_list,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
self.criterion = EdgeDetectionReweightedLosses()
self.model, self.optimizer = model, optimizer
# using cuda
if args.cuda:
self.model = DataParallelModel(self.model).cuda()
self.criterion = DataParallelCriterion(self.criterion).cuda()
# finetune from a trained model
if args.ft:
args.start_epoch = 0
checkpoint = torch.load(args.ft_resume)
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.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.ft_resume, checkpoint['epoch']))
# resuming checkpoint
if args.resume:
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.logger.info("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
# lr scheduler
self.scheduler = utils.LR_Scheduler(args.lr_scheduler,
args.lr,
args.epochs,
len(self.trainloader),
logger=self.logger,
lr_step=args.lr_step)
