def get_coco_data_loader(coco_image_dir, batch_size): transforms = DetectionPresetEval() coco_test_data = get_coco(root=coco_image_dir, image_set="val", transforms=transforms) coco_test_data_loader = DataLoader(coco_test_data, batch_size=batch_size, shuffle=False, collate_fn=collate_fn) return coco_test_data_loader
def __init__(self, data_dir='auto', split='train', use_crowded=False, return_crowded=False, return_area=False): if not _availabel: raise ValueError( 'Please install pycocotools \n' 'pip install -e \'git+https://github.com/pdollar/coco.git' '#egg=pycocotools&subdirectory=PythonAPI\'') self.use_crowded = use_crowded self.return_crowded = return_crowded self.return_area = return_area if split in ['val', 'minival', 'valminusminival']: img_split = 'val' else: img_split = 'train' if data_dir == 'auto': data_dir = get_coco(split, img_split) self.img_root = os.path.join(data_dir, 'images', '{}2014'.format(img_split)) anno_fn = os.path.join(data_dir, 'annotations', 'instances_{}2014.json'.format(split)) self.data_dir = data_dir anno = json.load(open(anno_fn, 'r')) self.img_props = dict() for img in anno['images']: self.img_props[img['id']] = img self.ids = list(self.img_props.keys()) cats = anno['categories'] self.cat_ids = [cat['id'] for cat in cats] self.anns = dict() self.imgToAnns = defaultdict(list) for ann in anno['annotations']: self.imgToAnns[ann['image_id']].append(ann) self.anns[ann['id']] = ann
def get_dataset(name, image_set, transform, data_path): paths = { "coco": (data_path, get_coco, 21), } DATA_DIR, _, num_classes = paths[name] ''' DATA_DIR:图片、标注存储根目录 coco_fabric_dataset:布匹数据、标注具体存储位置 ''' coco_fabric_dataset = { "img_dir": "/Users/liumingwei/machine-learning/tianchi/gd_fabric/experiment/coco/images/train", "ann_file": "/Users/liumingwei/machine-learning/tianchi/gd_fabric/experiment/coco/annotations/instances_train.json" } datesets = get_coco(DATA_DIR, image_set=image_set, data_set=coco_fabric_dataset, transforms=transform) return datesets, num_classes
def train_RCNN(model, path2data, path2json, weight_path=None): # train on the GPU or on the CPU, if a GPU is not available device = torch.device( 'cuda') if torch.cuda.is_available() else torch.device('cpu') #device = torch.device('cpu') # see if pretrained weights are available load_pretrained = False if weight_path is not None: load_pretrained = True # get coco style dataset dataset = coco_utils.get_coco(path2data, path2json, T.ToTensor()) # split the dataset in train and test set indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-1]) dataset_test = torch.utils.data.Subset(dataset, indices[-1:]) # define training and validation data loaders(use num_workers for multi-gpu) data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, collate_fn=myutils.collate_fn) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, collate_fn=myutils.collate_fn) if torch.cuda.device_count() > 1: print("Using", torch.cuda.device_count(), "GPUs") # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs model = nn.DataParallel(model) # move model to the right device model.to(device) # construct an optimizer params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.005, momentum=0.9, weight_decay=0.0005) # and a learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) # load the dataset in case of pretrained weights start_epoch = 0 if load_pretrained: checkpoint = torch.load(weight_path, map_location=device) model.load_state_dict(checkpoint['model_state_dict']) optimizer.load_state_dict(checkpoint['optimizer_state_dict']) start_epoch = checkpoint['epoch'] for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations train_one_epoch(model, optimizer, data_loader, device, epoch + start_epoch, print_freq=100) # update the learning rate #lr_scheduler.step() # evaluate on the test dataset # Find a way around the broken pytorch nograd keypoint evaluation # evaluate(model, data_loader_test, device=device) # save weights when done torch.save( { 'epoch': num_epochs + start_epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), }, save_weights_to)
def main(): #torch.manual_seed(42) # ------------ # args # ------------ parser = ArgumentParser() # Learning parameters parser.add_argument('--auto_lr', type=U.str2bool, default=False,help="Auto lr finder") parser.add_argument('--learning_rate', type=float, default=10e-4) parser.add_argument('--scheduler', type=U.str2bool, default=False) parser.add_argument('--wd', type=float, default=2e-4) parser.add_argument('--moment', type=float, default=0.9) parser.add_argument('--batch_size', default=5, type=int) parser.add_argument('--n_epochs', default=10, type=int) parser.add_argument('--iter_every', default=1, type=int,help="Accumulate compute graph for iter_size step") parser.add_argument('--benchmark', default=False, type=U.str2bool, help="enable or disable backends.cudnn") # Model and eval parser.add_argument('--model', default='FCN', type=str,help="FCN or DLV3 model") parser.add_argument('--pretrained', default=False, type=U.str2bool,help="Use pretrained pytorch model") parser.add_argument('--eval_angle', default=True, type=U.str2bool,help=\ "If true, it'll eval the model with different angle input size") # Data augmentation parser.add_argument('--rotate', default=False, type=U.str2bool,help="Use random rotation as data augmentation") parser.add_argument('--pi_rotate', default=True, type=U.str2bool,help="Use only pi/2 rotation angle") parser.add_argument('--p_rotate', default=0.25, type=float,help="Probability of rotating the image during the training") parser.add_argument('--scale', default=True, type=U.str2bool,help="Use scale as data augmentation") parser.add_argument('--landcover', default=False, type=U.str2bool,\ help="Use Landcover dataset instead of VOC and COCO") parser.add_argument('--size_img', default=520, type=int,help="Size of input images") parser.add_argument('--size_crop', default=480, type=int,help="Size of crop image during training") parser.add_argument('--angle_max', default=360, type=int,help="Angle max for data augmentation") # Dataloader and gpu parser.add_argument('--nw', default=0, type=int,help="Num workers for the data loader") parser.add_argument('--pm', default=True, type=U.str2bool,help="Pin memory for the dataloader") parser.add_argument('--gpu', default=0, type=int,help="Wich gpu to select for training") # Datasets parser.add_argument('--split', default=False, type=U.str2bool, help="Split the dataset") parser.add_argument('--split_ratio', default=0.3, type=float, help="Amount of data we used for training") parser.add_argument('--dataroot_voc', default='/data/voc2012', type=str) parser.add_argument('--dataroot_sbd', default='/data/sbd', type=str) parser.add_argument('--dataroot_landcover', default='/share/DEEPLEARNING/datasets/landcover', type=str) # Save parameters parser.add_argument('--model_name', type=str,help="what name to use for saving") parser.add_argument('--save_dir', default='/data/save_model', type=str) parser.add_argument('--save_all_ep', default=False, type=U.str2bool,help=\ "If true it'll save the model every epoch in save_dir") parser.add_argument('--save_best', default=False, type=U.str2bool,help="If true will only save the best epoch model") args = parser.parse_args() # ------------ # device # ------------ device = torch.device("cuda:"+str(args.gpu) if torch.cuda.is_available() else "cpu") print("device used:",device) # ------------ # data # ------------ if args.size_img < args.size_crop: raise Exception('Cannot have size of input images less than size of crop') size_img = (args.size_img,args.size_img) size_crop = (args.size_crop,args.size_crop) if not args.landcover: train_dataset_VOC = mdset.VOCSegmentation(args.dataroot_voc,year='2012', image_set='train', \ download=True,rotate=args.rotate,size_img=size_img,size_crop=size_crop) test_dataset = mdset.VOCSegmentation(args.dataroot_voc,year='2012', image_set='val', download=True) train_dataset_SBD = mdset.SBDataset(args.dataroot_sbd, image_set='train_noval',mode='segmentation',\ rotate=args.rotate,size_img=size_img,size_crop=size_crop) #COCO dataset if args.extra_coco: extra_COCO = cu.get_coco(args.dataroot_coco,'train',rotate=args.rotate,size_img=size_img,size_crop=size_crop) # Concatene dataset train_dataset = tud.ConcatDataset([train_dataset_VOC,train_dataset_SBD,extra_COCO]) else: train_dataset = tud.ConcatDataset([train_dataset_VOC,train_dataset_SBD]) num_classes = 21 else: print('Loading Landscape Dataset') train_dataset = mdset.LandscapeDataset(args.dataroot_landcover,image_set="trainval",\ rotate=args.rotate,pi_rotate=args.pi_rotate,p_rotate=args.p_rotate,size_img=size_img,size_crop=size_crop,angle_max=args.angle_max) test_dataset = mdset.LandscapeDataset(args.dataroot_landcover,image_set="test") print('Success load Landscape Dataset') num_classes = 4 split = args.split if split==True: train_dataset = U.split_dataset(train_dataset,args.split_ratio) # Print len datasets print("There is",len(train_dataset),"images for training and",len(test_dataset),"for validation") dataloader_train = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size,num_workers=args.nw,\ pin_memory=args.pm,shuffle=True,drop_last=True)#,collate_fn=U.my_collate) dataloader_val = torch.utils.data.DataLoader(test_dataset,num_workers=args.nw,pin_memory=args.pm,\ batch_size=args.batch_size) # ------------ # model # ------------ if args.model.upper()=='FCN': model = models.segmentation.fcn_resnet101(pretrained=args.pretrained,num_classes=num_classes) elif args.model.upper()=='DLV3': model = models.segmentation.deeplabv3_resnet101(pretrained=args.pretrained,num_classes=num_classes) else: raise Exception('model must be "FCN" or "DLV3"') #model.to(device) # ------------ # save # ------------ save_dir = U.create_save_directory(args.save_dir) print('model will be saved in',save_dir) U.save_hparams(args,save_dir) # ------------ # training # ------------ # Auto lr finding print(args) criterion = nn.CrossEntropyLoss(ignore_index=num_classes) # On ignore la classe border. torch.autograd.set_detect_anomaly(True) optimizer = torch.optim.SGD(model.parameters(),lr=args.learning_rate,momentum=args.moment,weight_decay=args.wd) ev.train_fully_supervised(model=model,n_epochs=args.n_epochs,train_loader=dataloader_train,val_loader=dataloader_val,\ criterion=criterion,optimizer=optimizer,save_folder=save_dir,scheduler=args.scheduler,auto_lr=args.auto_lr,\ model_name=args.model_name,benchmark=args.benchmark, save_best=args.save_best,save_all_ep=args.save_all_ep,\ device=device,num_classes=num_classes)
def main(): #torch.manual_seed(42) # ------------ # args # ------------ parser = ArgumentParser() parser.add_argument('--auto_lr', type=U.str2bool, default=False, help="Auto lr finder") parser.add_argument('--learning_rate', type=float, default=10e-4) parser.add_argument('--Loss', type=str, default='KL') parser.add_argument('--gamma', type=float, default=0.5, help="gamma balance the two losses") parser.add_argument('--scheduler', type=U.str2bool, default=True) parser.add_argument('--wd', type=float, default=2e-4) parser.add_argument('--moment', type=float, default=0.9) parser.add_argument('--batch_size', default=5, type=int) parser.add_argument('--iter_every', default=1, type=int, help="Accumulate compute graph for iter_size step") parser.add_argument('--n_epochs', default=10, type=int) parser.add_argument('--model', default='DLV3', type=str, help="FCN or DLV3 model") parser.add_argument('--pretrained', default=False, type=U.str2bool, help="Use pretrained pytorch model") parser.add_argument('--eval_angle', default=True, type=U.str2bool,help=\ "If true, it'll eval the model with different angle input size") parser.add_argument('--eval_every', default=30, type=int, help="Eval all input rotation angle every n step") parser.add_argument('--rotate', default=False, type=U.str2bool, help="Use random rotation as data augmentation") parser.add_argument('--angle_max', default=30, type=int, help="Max angle rotation of input image") parser.add_argument('--size_img', default=520, type=int, help="Size of input images") parser.add_argument('--size_crop', default=480, type=int, help="Size of crop image during training") parser.add_argument('--nw', default=0, type=int, help="Num workers for the data loader") parser.add_argument('--pm', default=True, type=U.str2bool, help="Pin memory for the dataloader") parser.add_argument('--gpu', default=0, type=int, help="Wich gpu to select for training") parser.add_argument( '--rot_cpu', default=False, type=U.str2bool, help="Apply rotation on the cpu (Help to use less gpu memory)") parser.add_argument('--benchmark', default=False, type=U.str2bool, help="enable or disable backends.cudnn") parser.add_argument('--split', default=True, type=U.str2bool, help="Split the dataset") parser.add_argument('--split_ratio', default=0.3, type=float, help="Amount of data we used for training") parser.add_argument('--extra_coco', default=False, type=U.str2bool,\ help="Use coco dataset as extra annotation for fully supervised training") parser.add_argument( '--multi_task', default=False, type=U.str2bool, help="Multi task training (same data for equiv and sup)") parser.add_argument('--dataroot_voc', default='/share/DEEPLEARNING/datasets/voc2012', type=str) parser.add_argument('--dataroot_sbd', default='/share/DEEPLEARNING/datasets/sbd', type=str) parser.add_argument('--dataroot_coco', default='/share/DEEPLEARNING/datasets/coco', type=str) parser.add_argument('--model_name', type=str, help="what name to use for saving") parser.add_argument('--save_dir', default='/data/save_model', type=str) parser.add_argument('--save_all_ep', default=False, type=U.str2bool,help=\ "If true it'll save the model every epoch in save_dir") parser.add_argument('--save_best', default=False, type=U.str2bool, help="If true will only save the best epoch model") parser.add_argument('--load_last_model', default=False, type=U.str2bool, help="If it will load the last model saved with\ This parameters." ) args = parser.parse_args() # ------------ # device # ------------ device = torch.device( "cuda:" + str(args.gpu) if torch.cuda.is_available() else "cpu") print("device used:", device) # ------------ # model # ------------ # ------------ # data # ------------ if args.size_img < args.size_crop: raise Exception( 'Cannot have size of input images less than size of crop') size_img = (args.size_img, args.size_img) size_crop = (args.size_crop, args.size_crop) train_dataset_VOC = mdset.VOCSegmentation(args.dataroot_voc,year='2012', image_set='train', \ download=True,rotate=args.rotate,size_img=size_img,size_crop=size_crop) val_dataset_VOC = mdset.VOCSegmentation(args.dataroot_voc, year='2012', image_set='val', download=True) train_dataset_SBD = mdset.SBDataset(args.dataroot_sbd, image_set='train_noval',mode='segmentation',\ rotate=args.rotate,size_img=size_img,size_crop=size_crop) #COCO dataset if args.extra_coco: extra_COCO = cu.get_coco(args.dataroot_coco, 'train', rotate=args.rotate, size_img=size_img, size_crop=size_crop) # Concatene dataset train_dataset_unsup = tud.ConcatDataset( [train_dataset_VOC, train_dataset_SBD]) # Split dataset split = args.split if split == True: train_dataset_sup = U.split_dataset(train_dataset_unsup, args.split_ratio) else: train_dataset_sup = train_dataset_unsup # Multi task ? if args.multi_task: train_dataset_unsup = train_dataset_sup # If extra coco concatene all dataset for unsupervised training if args.extra_coco: train_dataset_unsup = tud.ConcatDataset( [train_dataset_VOC, train_dataset_SBD, extra_COCO]) # Print len datasets print("There is",len(train_dataset_sup),"images for supervised training",len(train_dataset_unsup),\ "for equivariance loss and",len(val_dataset_VOC),"for validation") dataloader_train_sup = torch.utils.data.DataLoader(train_dataset_sup, batch_size=args.batch_size,num_workers=args.nw,\ pin_memory=args.pm,shuffle=True,drop_last=True) dataloader_val = torch.utils.data.DataLoader(val_dataset_VOC,num_workers=args.nw,pin_memory=args.pm,\ batch_size=args.batch_size) # --------- # Load model # --------- if args.load_last_model: model,save_dir = fbm.load_best_model(save_dir=args.save_dir,model_name=args.model_name,split=args.split,\ split_ratio=args.split_ratio,batch_size =args.batch_size,rotate=args.rotate) print("Training will continue from this file.", save_dir) else: save_dir = U.create_save_directory( args.save_dir) # Create a new save directory if args.model.upper() == 'FCN': model = models.segmentation.fcn_resnet101( pretrained=args.pretrained) elif args.model.upper() == 'DLV3': print('DEEPLAB MODEL') model = models.segmentation.deeplabv3_resnet101( pretrained=args.pretrained) else: raise Exception('model must be "FCN" or "DLV3"') model.to(device) # ------------ # save # ------------ print('model will be saved in', save_dir) U.save_hparams(args, save_dir) # ------------ # training # ------------ # Auto lr finding #if args.auto_lr==True: criterion_supervised = nn.CrossEntropyLoss( ignore_index=21) # On ignore la classe border. optimizer = torch.optim.SGD(model.parameters(), lr=args.learning_rate, momentum=args.moment, weight_decay=args.wd) ev.train_rot_equiv(model,args.n_epochs,dataloader_train_sup,train_dataset_unsup,dataloader_val,criterion_supervised,optimizer,\ scheduler=args.scheduler,Loss=args.Loss,gamma=args.gamma,batch_size=args.batch_size,iter_every=args.iter_every,save_folder=save_dir,\ model_name=args.model_name,benchmark=args.benchmark,angle_max=args.angle_max,size_img=args.size_img,\ eval_every=args.eval_every,save_all_ep=args.save_all_ep,dataroot_voc=args.dataroot_voc,save_best=args.save_best\ ,rot_cpu=args.rot_cpu,device=device) # Final evaluation """