def __init__(self, args): self.args = 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} trainset = get_segmentation_dataset(args.dataset, split=args.train_split, mode='train', root=args.data_folder, **data_kwargs) testset = get_segmentation_dataset(args.dataset, split='val', mode ='val', root=args.data_folder, **data_kwargs) # dataloader kwargs = {'num_workers': args.workers, 'pin_memory': True} 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 = args.num_classes self.best_pred = 0.0 # model if args.model == "unet": from code.models import UNet model = UNet(n_channels=3, n_classes=args.num_classes, bilinear=False) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) else: model = fpn.get_fpn(nclass=args.num_classes, backbone=args.backbone, pretrained=False) # 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}) # optimizer = torch.optim.SGD(params_list, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay) optimizer = torch.optim.Adam(params_list, lr=args.lr, weight_decay=args.weight_decay) # criterions self.criterion = torch.nn.CrossEntropyLoss(ignore_index=args.ignore_label) self.model, self.optimizer = model, optimizer # using cuda if args.cuda: self.model = self.model.cuda() self.criterion = self.criterion.cuda() # resuming checkpoint if args.weight is not None: if not os.path.isfile(args.weight): raise RuntimeError("=> no checkpoint found at '{}'" .format(args.weight)) checkpoint = torch.load(args.weight, map_location='cuda:0') checkpoint['state_dict'] = OrderedDict([(k[5:], v) if 'base' in k else (k, v) for k, v in checkpoint['state_dict'].items()]) args.start_epoch = checkpoint['epoch'] if args.cuda: self.model.load_state_dict(checkpoint['state_dict'], strict=False) else: self.model.load_state_dict(checkpoint['state_dict']) self.optimizer.load_state_dict(checkpoint['optimizer']) self.best_pred = checkpoint['best_pred'] print("=> loaded checkpoint '{}' (epoch {})" \ .format(args.weight, checkpoint['epoch'])) # clear start epoch if fine-tuning self.scheduler = LR_Scheduler(args.lr_scheduler, args.lr, \ args.epochs, len(self.trainloader), warmup_epochs=5)
def analyze(self): from code.datasets import get_segmentation_dataset input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) data_kwargs = { 'transform': input_transform, 'base_size': args.base_size, 'crop_size': args.crop_size } testset = get_segmentation_dataset(args.dataset, split='val', mode='testval', root=args.data_folder, **data_kwargs) loader_kwargs = { 'num_workers': args.workers, 'pin_memory': True } if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False) self.pruner.ana(ana_eval_fn, args=(test_data, self.criterion))
def build_data(args, subset_len=None, sample_method='random'): from code.datasets import get_segmentation_dataset input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) data_kwargs = { 'transform': input_transform, 'base_size': args.base_size, 'crop_size': args.crop_size } if args.eval: testset = get_segmentation_dataset(args.dataset, split='val', mode='testval', root=args.data_folder, **data_kwargs) loader_kwargs = { 'num_workers': args.workers, 'pin_memory': True } if args.cuda else {} if subset_len: assert subset_len <= len(testset) if sample_method == 'random': testset = torch.utils.data.Subset( testset, random.sample(range(0, len(test_data)), subset_len)) else: testset = torch.utils.data.Subset(testset, list(range(subset_len))) #dataloader test_data = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False) return test_data
last_snapshot_name = None best_snapshot = None if (args.eval_pruned == True): from code.datasets import get_segmentation_dataset input_transform = transform.Compose([ transform.ToTensor(), transform.Normalize([.485, .456, .406], [.229, .224, .225]) ]) data_kwargs = { 'transform': input_transform, 'base_size': args.base_size, 'crop_size': args.crop_size } testset = get_segmentation_dataset(args.dataset, split='val', mode='testval', root=args.data_folder, **data_kwargs) loader_kwargs = { 'num_workers': args.workers, 'pin_memory': True } if args.cuda else {} test_data = data.DataLoader(testset, batch_size=args.batch_size, drop_last=False, shuffle=False) evaluate(test_data, trainer.model, trainer.criterion) if (args.prune == True): if (args.prune_model_py is None and args.pruned_weights is None): trainer.analyze()