def main(): fixed_seed(42) if args.mode == 'train': model = get_model_instance_segmentation(args.num_classes).to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = optim.SGD(params, lr=1e-3, momentum=0.9, weight_decay=1e-4) for epoch in range(args.epochs): train_fn(model, train_loader, optimizer, device, epoch) torch.save(model.state_dict(), f'./weight/test{epoch}.pth') else: model = get_model_instance_segmentation(args.num_classes).to(device) model.load_state_dict(torch.load('./weight/test1.pth')) model.eval() test_fn(model, test_loader, class_nums, device)
def main(): fixed_seed(42) model = get_model_instance_segmentation(args.num_classes).to(device) model.load_state_dict(torch.load('./weight/test1.pth')) model.eval() test_fn(model, test_loader, class_nums, device)
def predict(image): model = get_model_instance_segmentation(2) model.load_state_dict( torch.load('best_model', map_location=torch.device('cpu'))) model.eval() image = Image.open(image).convert('RGB') img = loader(image) with torch.no_grad(): output = model(img[None, ...])[0] display_instances(image, output['boxes'], output['labels'], class_names, output['scores'])
def run_train(): fixed_seed(42) if not os.path.exists('./weight/'): print('Not exists ./weight/ making an weight folder...') os.mkdir('./weight/') model = get_model_instance_segmentation(args.num_classes).to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = optim.SGD(params, lr=1e-3, momentum=0.9, weight_decay=1e-4) for epoch in range(args.epochs): train_fn(model, train_loader, optimizer, device, epoch) torch.save(model.state_dict(), f'./weight/test{epoch}.pth')
def main(): parser = argparse.ArgumentParser(description='Mask R-CNN') parser.add_argument( '--image_path', default= '/home/hatsunemiku/dev/mask-rcnn/data/food201/segmented_test/pizza/309892.jpg', type=str, help='image path') args = parser.parse_args() model = get_model_instance_segmentation(num_classes=2) model.load_state_dict(torch.load('food201_model.pkl')) model.eval() detect_and_color_splash(model, args.image_path, dataset_name='food201')
def train(config): # create our own dataset and its data_loader tr_dt = FashionDataset(config, get_transform(train = True)) tr_data_loader = DataLoader( tr_dt, config.batch_size , shuffle = True, num_workers = 8, collate_fn = lambda x: tuple(zip(*x)) ) # save the weights # weight_file = osp.join(config.save_dir, 'weights') # check them whether exists or not if not os.path.exists(config.save_dir): os.makedirs(config.save_dir) # there are 46 classes in total num_classes = 46 + 1 # create model instance model = get_model_instance_segmentation(num_classes) #set model to device model.to(device) # for optim params = [p for p in model.parameters() if p.requires_grad] optim = torch.optim.SGD(params, lr = 0.001, momentum=0.9, weight_decay=0.0005) lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=5, gamma=0.1) for epoch in range(config.num_epochs): train_one_epoch(model, optim, tr_data_loader, device, epoch, print_freq=config.rep_intv) # updt the learning rate lr_scheduler.step() w1 = osp.join(config.save_dir , 'weights') wfile = osp.join(w1, '{}_model.bin'.format(str(epoch))) torch.save(model.state_dict(), wfile)
def main(): # 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') print(f'Device: {device}') # our dataset has two classes only - background and person num_classes = 2 # use our dataset and defined transformations dataset = PennFudanDataset('PennFudanPed', get_transform(train=True)) dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False)) # split the dataset in train and test set indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-50]) dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:]) # define training and validation data loaders data_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True, num_workers=1, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=1, collate_fn=utils.collate_fn) # get the model using our helper function model = get_model_instance_segmentation(num_classes) # 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) # let's train it for 10 epochs num_epochs = 10 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) # update the learning rate lr_scheduler.step() # evaluate on the test dataset evaluate(model, data_loader_test, device=device) print("That's it!")
# split the dataset in train and test set #indices = torch.randperm(len(dataset)).tolist() #dataset = torch.utils.data.Subset(dataset, indices) #indices_test = torch.randperm(len(dataset_test)).tolist() #dataset_test = torch.utils.data.Subset(dataset_test, indices_test) # define training and validation data loaders data_loader = torch.utils.data.DataLoader( dataset, batch_size=4, shuffle=True, num_workers=8, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, num_workers=4, collate_fn=utils.collate_fn) # get the model using our helper function mask_rcnn = get_model_instance_segmentation(num_classes, image_mean, image_std, stats=True) read_param = False if read_param: mask_rcnn.load_state_dict(torch.load("trained_param_6/epoch_0050.param")) print("Loaded weights") # move model to the right device mask_rcnn.to(device) # construct an optimizer params = [p for p in mask_rcnn.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=40,
if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument("-b", "--backbone", type=str, choices=['resnet', 'mn2'], default='resnet', help="backbone type") parser.add_argument("-m", "--modelpath", default='./checkpoint/inb_resnet_9', help="model file path") parser.add_argument("-d", "--datapath", default='./benchmark', help="data path") args = parser.parse_args() dataset = INBDataset(args.datapath, M.get_transform(train=False)) dataset_test = dataset num_classes = 2 model = M.get_model_instance_segmentation(args.backbone, num_classes) checkpoint = torch.load(args.modelpath) # model.load_state_dict(checkpoint['model_state_dict']) model.load_state_dict(checkpoint) device = torch.device('cuda') if torch.cuda.is_available()\ else torch.device('cpu') model.to(device) stat(model, dataset_test, device)
batch_size=2, shuffle=True, num_workers=2, collate_fn=dataset.collate_fn, ) dataloader = {"train": dataloader_train, "val": dataloader_test} if checkpoint is None: # train on the GPU or on the CPU, if a GPU is not available os.environ["CUDA_VISIBLE_DEVICES"] = "1" device = (torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")) print(device) with torch.cuda.device(0): # load a pre-trained model model = get_model_instance_segmentation(21) optimizer_ft = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=0.0005) # and a learning rate scheduler milestones = [10, 23, 33, 43] exp_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer_ft, milestones, gamma=0.1, last_epoch=-1) # exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer_ft, # step_size=10, # gamma=0.1)
# our dataset has three classes only - background, non-damaged, and damaged num_classes = 6 # 3 or 6 dataset_test = Dataset("./datasets/Eureka_infer/102/", "./datasets/Eureka_infer/102_labels/", get_transform(train=False), readsave=False) data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=2, collate_fn=utils.collate_fn) mask_rcnn = get_model_instance_segmentation(num_classes, image_mean=None, image_std=None, stats=False) mask_rcnn.load_state_dict( torch.load("trained_param_eureka_aug_mult/epoch_0021.param")) print("loaded weights") # move model to the right device mask_rcnn.to(device) mask_rcnn.eval() for image, target in dataset_test: pred = mask_rcnn(image.unsqueeze(0).to(device))[0] boxes = pred['boxes'].to("cpu").detach().numpy()
def main(): # 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') print(device) # our dataset has two classes only - background and person num_classes = 6 # use our dataset and defined transformations MASK_MAP = {"title": 1, "rating": 2, "ratings": 3, "price": 4, "description": 5} dataset = LabelboxWebpageDataset("../data/product_images", "../data/labelbox-export.862Z.json", MASK_MAP, get_transform(True)) dataset_test = LabelboxWebpageDataset("../data/product_images", "../data/labelbox-export.862Z.json", MASK_MAP, get_transform(False)) # split the dataset in train and test set indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-5]) dataset_test = torch.utils.data.Subset(dataset_test, indices[-5:]) # define training and validatio n data loaders data_loader = torch.utils.data.DataLoader( dataset, batch_size=1, shuffle=True, num_workers=0, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, num_workers=0, collate_fn=utils.collate_fn) # get the model using our helper function model = get_model_instance_segmentation(num_classes) # 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) optimizer = torch.optim.Adam(params, lr=0.0005) # and a learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) # let's train it for 10 epochs num_epochs = 10 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) # update the learning rate lr_scheduler.step() # evaluate on the test dataset # Evaluation broken if epoch == num_epochs - 1: model.eval() # cpu_device = torch.device("cpu") # model.to(cpu_device) # for images, targets in data_loader_test: # # images = list(img.to(cpu_device) for img in images) # # outputs = model(images) # # print(outputs) # # print(targets) evaluate(model, data_loader_test, device=device)
def main(): # 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') # our dataset has two classes only - background and person num_classes = 2 # get the model using our helper function model = get_model_instance_segmentation(num_classes) # load pretrain_dict pretrain_dict = torch.load( os.path.join("C:\\zhulei\\maskRcnn\\models", "_epoch-9.pth")) model.load_state_dict(pretrain_dict) # move model to the right device model.to(device) # use our dataset and defined transformations dataset_test = PennFudanDataset('C:\\zhulei\\maskRcnn\\data\\test', get_transform(train=False)) data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=1, collate_fn=utils.collate_fn) n_threads = torch.get_num_threads() # FIXME remove this and make paste_masks_in_image run on the GPU torch.set_num_threads(1) cpu_device = torch.device("cpu") model.eval() metric_logger = utils.MetricLogger(delimiter=" ") header = 'Test:' coco = get_coco_api_from_dataset(data_loader_test.dataset) iou_types = _get_iou_types(model) coco_evaluator = CocoEvaluator(coco, iou_types) for image, targets in metric_logger.log_every(data_loader_test, 100, header): image = list(img.to(device) for img in image) targets = [{k: v.to(device) for k, v in t.items()} for t in targets] torch.cuda.synchronize() model_time = time.time() outputs = model(image) instance_segmentation_api(image[0], outputs) # 可视化 # for img in image: # Image.fromarray((img.mul(255).permute(1, 2, 0).byte().cpu().numpy())[0]) # print(outputs[0]['masks'].shape) # for i in range(99): # result = Image.fromarray(outputs[0]['masks'][i, 0].mul(255).byte().cpu().numpy()) # result.show() outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs] model_time = time.time() - model_time res = { target["image_id"].item(): output for target, output in zip(targets, outputs) } evaluator_time = time.time() coco_evaluator.update(res) evaluator_time = time.time() - evaluator_time metric_logger.update(model_time=model_time, evaluator_time=evaluator_time) # gather the stats from all processes metric_logger.synchronize_between_processes() print("Averaged stats:", metric_logger) coco_evaluator.synchronize_between_processes() # accumulate predictions from all images coco_evaluator.accumulate() coco_evaluator.summarize() torch.set_num_threads(n_threads)
def get_coloured_mask(mask, pred_cls): r = np.zeros_like(mask).astype(np.uint8) g = np.zeros_like(mask).astype(np.uint8) b = np.zeros_like(mask).astype(np.uint8) colours = [[0, 0, 0],[0, 255, 0],[0, 255, 255],[255, 255, 0],[80, 70, 180],[180, 40, 250],[245, 145, 50],[70, 150, 250],[50, 190, 190]] r[mask == 1], g[mask == 1], b[mask == 1] = colours[common.CLASS_NAMES.index(pred_cls)] coloured_mask = np.stack([r, g, b], axis=2) return coloured_mask if __name__ == '__main__': # load images img_paths = glob.glob(common.TRAIN_DATA_IMGS) # load model model = model.get_model_instance_segmentation(common.NUM_CLASSES) model, device = common.setup_device(model) model.load_state_dict(torch.load("model.pth", map_location=device)) model.eval() confidence = 0.5 for idx in range(len(img_paths)): # Prediction img_path = img_paths[idx] img = Image.open(img_path) transform = T.Compose([T.ToTensor()]) img = torchvision.transforms.functional.to_tensor(img) pred = model([img.to(device)]) pred_score = list(pred[0]['scores'].detach().cpu().numpy()) pred_t = [pred_score.index(x) for x in pred_score if x>confidence]
def generate_mask(img_base, img_name): if torch.cuda.is_available(): device = torch.device('cuda') else: device = torch.device('cpu') num_classes = 2 image_mean = [ 0.34616187074865956, 0.34616187074865956, 0.34616187074865956 ] image_std = [0.10754412766560431, 0.10754412766560431, 0.10754412766560431] mask_rcnn = get_model_instance_segmentation(num_classes, image_mean, image_std, stats=True) mask_rcnn.to(device) mask_rcnn.eval() model_path = "models/epoch_0099.param" mask_rcnn.load_state_dict(torch.load(model_path)) # tif files and png files location img_path = img_base + "/split_pngs/" tif_path = img_base + "/split_tifs/" # tmp folder has the vector files created for each craters # create if it doesn't exist tmp_folder = img_base + "/split_tifs/tmp/" if not os.path.exists(tmp_folder): os.makedirs(tmp_folder) pred_folder = img_base + "/split_tifs/shape_files/" if not os.path.exists(pred_folder): os.makedirs(pred_folder) print(img_name) # file name without extension file_name = img_name.split(".")[0] individual_folder = tmp_folder + file_name + "/" if not os.path.exists(individual_folder): os.makedirs(individual_folder) image = Image.open(img_path + img_name).convert("RGB") image = np.array(image) image = torch.from_numpy(image / 255.0).float() image = image.permute((2, 0, 1)) pred = mask_rcnn(image.unsqueeze(0).to(device))[0] #from visualize import display_instances boxes_ = pred["boxes"].cpu().detach().numpy().astype(int) boxes = np.empty_like(boxes_) boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3] = boxes_[:, 1], boxes_[:, 0], boxes_[:, 3], boxes_[:, 2] labels = pred["labels"].cpu().detach().numpy() scores = pred["scores"].cpu().detach().numpy() masks = pred["masks"] indices = scores > 0.6 # Show Scores boxes = boxes[indices] labels = labels[indices] scores = scores[indices] masks = masks[indices].squeeze(1) masks = (masks.permute( (1, 2, 0)).cpu().detach().numpy() > 0.5).astype(np.uint8) # Show image predictions and save as npy files image = image.permute((1, 2, 0)).cpu().detach().numpy() * 255 #display_instances(image, boxes, masks, labels, class_names=["background", "crater"], scores=scores) print(masks.shape) # Save masks as tif files crater_prediction_count = masks.shape[2] for i in range(crater_prediction_count): # create single crater tif file tmp_tif = individual_folder + file_name + "_" + str(i) + ".tif" os.system("cp " + tif_path + file_name + ".tif " + tmp_tif) dataset = gdal.Open(tmp_tif, gdal.GA_Update) band = dataset.GetRasterBand(1) new_mask = masks[:, :, i] new_mask = new_mask band.SetNoDataValue(0) band.WriteArray(new_mask) dataset.FlushCache() del dataset # convert the tif file to a vector data gdal polygonize tmp_shp = individual_folder + file_name + "_" + str(i) + ".shp" os.system( """gdal_polygonize.py {} {} -b 1 -f "ESRI Shapefile" 1 DN """. format(tmp_tif, tmp_shp)) if crater_prediction_count > 0: os.system("ogrmerge.py -single -overwrite_ds -o {} {}*.shp".format( pred_folder + file_name + ".shp", individual_folder)) os.system("rm -rf " + individual_folder + "*")
def test(config): # test_dt = FashionDataset(config, transforms= None) sample_df = pd.read_csv(config.sample_path) ################################################################################ # create the model instance num_classes = 46 + 1 model_test = get_model_instance_segmentation(num_classes) #load the training weights # load_path =osp.join(config.save_dir, '9_weights') load_path = osp.join(config.save_dir, 'weights') model_test.load_state_dict(torch.load(osp.join(load_path, '9_model.bin'))) # send the test model to gpu model_test.to(device) for param in model_test.parameters(): param.requires_grad = False model_test.eval() # for submission sub_list = [] missing_count = 0 for i, row in tqdm(sample_df.iterrows(), total=len(sample_df)): ###modify########################################################## # import the image img_path = osp.join(config.test_dir, sample_df['ImageId'][i]) img = Image.open(img_path).convert('RGB') img = img.resize((config.width, config.height), resample=Image.BILINEAR) # convert the img as tensor img = F.to_tensor(img) #####modify############################################################# pred = model_test([img.to(device)])[0] masks = np.zeros((512, 512, len(pred['masks']))) for j, m in enumerate(pred['masks']): res = transforms.ToPILImage()(m.permute(1, 2, 0).cpu().numpy()) res = np.asarray(res.resize((512, 512), resample=Image.BILINEAR)) masks[:, :, j] = (res[:, :] * 255. > 127).astype(np.uint8) labels = pred['labels'].cpu().numpy() scores = pred['scores'].cpu().numpy() # best_idx = 0 # print('the maximum scores is {}'.format(np.mean(scores))) # print('the current masks is {}'.format(masks)) for _scores in scores: if _scores > 0.8: best_idx += 1 if best_idx == 0: # print(masks.shape[-1]) sub_list.append([sample_df.loc[i, 'ImageId'], '1 1', 23]) missing_count += 1 continue if masks.shape[-1] > 0: masks = refine_masks(masks[:, :, :best_idx], labels[:best_idx]) for m, rle, label in masks: sub_list.append([sample_df.loc[i, 'ImageId'], rle, label]) else: sub_list.append([sample_df.loc[i, 'ImageId'], '1 1', 23]) missing_count += 1 submission_df = pd.DataFrame(sub_list, columns=sample_df.columns.values) print("Total image results: ", submission_df['ImageId'].nunique()) print("Missing Images: ", missing_count) submission_df = submission_df[submission_df.EncodedPixels.notnull()] for row in range(len(submission_df)): line = submission_df.iloc[row, :] submission_df.iloc[row, 1] = line['EncodedPixels'].replace('.0', '') # submission_df.head() submit_path = config.submit_path + r'submission.csv' print('submit_path: ', submit_path) submission_df.to_csv(submit_path, index=False) print('ok,finished')
def main(): parser = argparse.ArgumentParser(description='Mask R-CNN') parser.add_argument('--dataset', default='unimib', type=str, choices=['unimib', 'food201', 'chfood'], help='dataset name') args = parser.parse_args() # 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') if args.dataset == 'unimib': num_classes = 74 # use our dataset and defined transformations dataset = UNIMIBDataset( '/home/hatsunemiku/dev/mask-rcnn/data/UNIMIB2016', get_transform(train=True)) dataset_test = UNIMIBDataset( '/home/hatsunemiku/dev/mask-rcnn/data/UNIMIB2016', get_transform(train=False), False) elif args.dataset == 'food201': num_classes = 2 # use our dataset and defined transformations dataset = Food201Dataset( '/home/hatsunemiku/dev/mask-rcnn/data/food201', get_transform(train=True)) dataset_test = Food201Dataset( '/home/hatsunemiku/dev/mask-rcnn/data/food201', get_transform(train=False), False) elif args.dataset == 'chfood': num_classes = 24 # use our dataset and defined transformations dataset = CHFoodDataset('/home/hatsunemiku/dev/mask-rcnn/data/ch_food', get_transform(train=True)) dataset_test = CHFoodDataset( '/home/hatsunemiku/dev/mask-rcnn/data/ch_food', get_transform(train=False), False) else: raise Exception() # define training and validation data loaders data_loader = torch.utils.data.DataLoader(dataset, batch_size=2, shuffle=True, num_workers=0, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader(dataset_test, batch_size=1, shuffle=False, num_workers=0, collate_fn=utils.collate_fn) # get the model using our helper function model = get_model_instance_segmentation(num_classes) # 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.002, momentum=0.9, weight_decay=0.0005) # and a learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.1) # let's train it for 10 epochs num_epochs = 30 for epoch in range(num_epochs): # train for one epoch, printing every 10 iterations train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10) print('epoch %d, lr -----------> %.6f' % (epoch, optimizer.param_groups[0]['lr'])) # update the learning rate lr_scheduler.step() # evaluate on the test dataset evaluate(model, data_loader_test, device=device) torch.save(model.state_dict(), 'chfood_model.pkl')
indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices) indices_test = torch.randperm(len(dataset_test)).tolist() dataset_test = torch.utils.data.Subset(dataset_test, indices_test) # define training and validation data loaders data_loader = torch.utils.data.DataLoader( dataset, batch_size=4, shuffle=True, num_workers=4, collate_fn=utils.collate_fn) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, num_workers=4, collate_fn=utils.collate_fn) # get the model using our helper function mask_rcnn = get_model_instance_segmentation(num_classes, None, None, False) read_param = False if read_param: mask_rcnn.load_state_dict(torch.load("trained_param_eureka_aug_bin/epoch_0009.param")) # move model to the right device mask_rcnn.to(device) # construct an optimizer params = [p for p in mask_rcnn.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=0.01, momentum=0.9, weight_decay=0.00001) # and a learning rate scheduler lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10,
def main(): # mode argument args = argparse.ArgumentParser() args.add_argument("--num_classes", type=int, default=22) args.add_argument("--lr", type=int, default=0.005) args.add_argument("--cuda", type=bool, default=True) args.add_argument("--num_epochs", type=int, default=30) args.add_argument("--model_name", type=str, default="mask_1.pth") args.add_argument("--prediction_file", type=str, default="prediction") args.add_argument("--batch", type=int, default=16) args.add_argument("--mode", type=str, default="train") config = args.parse_args() num_classes = config.num_classes base_lr = config.lr cuda = config.cuda num_epochs = config.num_epochs model_name = config.model_name prediction_file = config.prediction_file batch = config.batch mode = config.mode # 도움 함수를 이용해 모델을 가져옵니다 new_model = model.get_model_instance_segmentation(num_classes) # 학습을 GPU로 진행하되 GPU가 가용하지 않으면 CPU로 합니다 device = torch.device('cuda') if cuda else torch.device('cpu') # 모델을 GPU나 CPU로 옮깁니다 new_model.to(device) if mode == 'train': # 데이터셋과 정의된 변환들을 사용합니다 '''dataset = CustomDataset(DATASET_PATH, dataloader.get_transform(train=True)) dataset_val = CustomDataset(DATASET_PATH, dataloader.get_transform(train=False)) # 데이터셋을 학습용과 테스트용으로 나눕니다(역자주: 여기서는 전체의 50개를 테스트에, 나머지를 학습에 사용합니다) indices = torch.randperm(len(dataset)).tolist() dataset = torch.utils.data.Subset(dataset, indices[:-10]) dataset_val = torch.utils.data.Subset(dataset_val, indices[-10:]) data_loader_val = torch.utils.data.DataLoader( dataset_val, batch_size=batch, shuffle=True,num_workers=4, collate_fn=dataloader.collate_fn) ''' dataset =dataloader.make_dataset(DATASET_PATH) # 데이터 로더를 학습용과 검증용으로 정의합니다 dataset_loader = torch.utils.data.DataLoader( dataset, batch_size=batch, shuffle=True, num_workers=4,collate_fn=dataloader.collate_fn) # 옵티마이저(Optimizer)를 만듭니다 params = [p for p in new_model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=base_lr, momentum=0.9, weight_decay=0.0005) # 학습률 스케쥴러를 만듭니다 lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,step_size=10,gamma=0.1) train(new_model, dataset_loader, device, num_epochs, optimizer=optimizer, lr_scheduler=lr_scheduler) elif mode == 'test' : dataset_test =dataloader.make_dataset(DATASET_PATH) data_loader_test = torch.utils.data.DataLoader( dataset_test, batch_size=1, shuffle=False, collate_fn=dataloader.collate_fn) load_model(model_name, new_model) test(new_model, data_loader_test, device, prediction_file) print("That's it!")
trainset = torch.utils.data.Subset(trainset, indeces[:-50]) testset = torch.utils.data.Subset(testset, indeces[-50:]) trainloader = torch.utils.data.DataLoader(trainset, batch_size=2, num_workers=4, shuffle=True, collate_fn=utils.collate_fn) testloader = torch.utils.data.DataLoader(testset, batch_size=1, num_workers=4, shuffle=True, collate_fn=utils.collate_fn) num_classes = 2 model = get_model_instance_segmentation(num_classes) model.to(device) params = [p for p in model.parameters() if p.requires_grad] optimizer = torch.optim.SGD(params, lr=5e-3, momentum=0.9, weight_decay=5e-4) lr_schedular = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1) epochs = 10 for epoch in range(epochs): train_one_epoch(model, optimizer, trainloader, device, epoch, print_freq=10) lr_schedular.step() evaluate(model, testloader, device)
def __init__(self): super(Ui_MainWindow, self).__init__() self.num_classes = 46 + 1 self.model_test = get_model_instance_segmentation(self.num_classes) self.model_path = './9_model.bin' self.pretrain_params = torch.load(self.model_path, map_location='cpu') for k in list(self.pretrain_params.keys()): if k.startswith('module.'): self.pretrain_params[ k[len('module.'):]] = self.pretrain_params[k] del self.pretrain_params[k] self.model_test.load_state_dict(self.pretrain_params) # send the test model to gpu self.model_test.to(device) for param in self.model_test.parameters(): param.requires_grad = False self.model_test.eval() #self.config, self.unparsed = get_config() self.image_name = None self.raw_image = None self.logo_size = (350, 350) self.mask_size = (60, 60) self.width = 512 self.height = 512 self.pred = None self.classes = None self.save_path = './images_pred/' with open('label_descriptions.json') as json_f: look_up_classes = json.load(json_f) self.categories = look_up_classes["categories"] self.attributes = look_up_classes["attributes"] if not os.path.exists('./images_pred/'): os.makedirs('./images_pred/') self.dic = { #'003b3e4579833df109d330773418fd65':{33:[182],31:[160,204],24:[],10:[108,141,155]}, '008324a52c5910d338807fb785ebaa0a': { 13: [], 10: [101, 108, 128, 142], 31: [158, 204], 28: [163], 23: [] }, 'wen': { 33: [183], 9: [91, 115, 136, 145], 29: [174], 32: [218], 31: [160] }, '3people': { 1: [0, 115, 145], 31: [157, 204], 24: [], 33: [182] }, #'00beb7542b35409dd316572e497e701c':{8:[65,115,118,229],1:[11,142],33:[],24:[]} } self.setupUi()