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)