def train(network_specs, training_params, image_path, save_path, ckpt_path, epoch=10): print('creating datapipe...') # create images DataPipeline datapipe = DataPipeline(image_path=image_path, training_params=training_params) print('creating network model...') # create model VAE model = UNet(network_specs=network_specs, datapipe=datapipe, training_params=training_params) # train the model # save_config is flexible print(''' ============= HERE WE GO ============= ''') model.train(save_path=save_path, ckpt_path=ckpt_path, epoch=epoch)
def train(network_specs, training_params, image_path, save_path, ckpt_path): print('creating datapipe...') # create images DataPipeline datapipe = DataPipeline(image_path=image_path, training_params=training_params) print('creating network model...') # create model VAE model = UNet(network_specs=network_specs, datapipe=datapipe, training_params=training_params, mode='evaluating') # train the model # save_config is flexible print(''' ============= HERE WE GO ============= ''') images, labels, preds = model.evaluate(ckpt_path=ckpt_path) for i in range(labels.shape[0]): # plt.subplots(figsize=[16,12]) for j in range(3): plt.subplot(2, 6, 2 * j + 1) plt.imshow(labels[i, :, :, j]) plt.subplot(2, 6, 2 * j + 2) plt.imshow(preds[i, :, :, j]) plt.subplot(2, 6, 7) plt.imshow(np.squeeze(images[i])) plt.show()
def train(epochs, batch_size, learning_rate): train_loader = torch.utils.data.DataLoader(SegThorDataset( "data", phase='train', transform=transforms.Compose([Rescale(0.25), Normalize(), ToTensor()]), target_transform=transforms.Compose([Rescale(0.25), ToTensor()])), batch_size=batch_size, shuffle=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = UNet().to(device) model.apply(weight_init) #optimizer = optim.Adam(model.parameters(), lr=learning_rate) #learning rate to 0.001 for initial stage optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.95) #optimizer = adabound.AdaBound(params = model.parameters(), lr = 0.001, final_lr = 0.1) for epoch in range(epochs): print('Epoch {}/{}'.format(epoch + 1, epochs)) print('-' * 10) running_loss = 0.0 loss_seg = np.zeros(5) for batch_idx, (train_data, labels) in enumerate(train_loader): train_data, labels = train_data.to( device, dtype=torch.float), labels.to(device, dtype=torch.uint8) print("train data size", train_data.size()) print("label size", labels.size()) optimizer.zero_grad() output = model(train_data) print("output: {} and taget: {}".format(output.size(), labels.size())) loss_label, loss = dice_loss(output, labels) loss.backward() optimizer.step() running_loss += loss.item() for i in range(4): loss_seg[i] += loss_label[i] print("Length: ", len(train_loader)) epoch_loss = running_loss / len(train_loader) epoch_loss_class = np.true_divide(loss_seg, len(train_loader)) print( "Dice per class: Background = {:.4f} Eusophagus = {:.4f} Heart = {:.4f} Trachea = {:.4f} Aorta = {:.4f}\n" .format(epoch_loss_class[0], epoch_loss_class[1], epoch_loss_class[2], epoch_loss_class[3], epoch_loss_class[4])) print("Total Dice Loss: {:.4f}\n".format(epoch_loss)) os.makedirs("models", exist_ok=True) torch.save(model, "models/model.pt")
def main(): FLAGS = parser.parse_args() # Calculate the predictions for all the images in the input_img_dir. for img_name in os.listdir(FLAGS.input_img_dir): if img_name.endswith('.png'): original_img = load_img( FLAGS.input_img_dir + '/' + img_name) # load_img function exists in file utils.py # Resizing image because of the small memory size input_img = tl.prepro.imresize(original_img, [FLAGS.img_size, FLAGS.img_size]) input_img = np.reshape(input_img, [1, FLAGS.img_size, FLAGS.img_size, 3]) unet = UNet(FLAGS.img_size) # The output is an array of the size(img_size * img_size, 1) prediction = unet.predict(input_img, FLAGS.model_save_dir) # Saving the image given the probabilities # save_img fnuction exists in file utils.py save_img( prediction, original_img, FLAGS.img_size, FLAGS.input_img_dir + '/' + img_name.split('.')[0] + '_pred.png')
def train(): ex = wandb.init(project="PQRST-segmentation") ex.config.setdefaults(wandb_config) logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') logging.info(f'Using device {device}') net = UNet(in_ch=1, out_ch=4) net.to(device) try: train_model(net=net, device=device, batch_size=wandb.config.batch_size, lr=wandb.config.lr, epochs=wandb.config.epochs) except KeyboardInterrupt: try: save = input("save?(y/n)") if save == "y": torch.save(net.state_dict(), 'net_params.pkl') sys.exit(0) except SystemExit: os._exit(0)
def main(): train_root_dir = '/content/drive/My Drive/DDSM/train/CBIS-DDSM' test_root_dir = '/content/drive/My Drive/DDSM/test/CBIS-DDSM' path_weights = '/content/drive/My Drive/Cv/weights' batch_size = 3 valid_size = 0.2 nb_epochs = 20 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # data loaders loaders = dataloaders(train_root_dir, combined_transform, batch_size, valid_size) model = UNet(in_channels=3, out_channels=1) model.to(device) optimizer = optim.Adam(model.parameters(), lr=0.01) exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=4, gamma=0.3) model = train(model, optimizer, exp_lr_scheduler, loaders, nb_epochs, device, path_weights) # from torchsummary import summary # # summary(model, input_size=(3, 224, 224)) # test_transform = transforms.Compose([ # transforms.ToTensor(), # transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) # ]) test_loader = DataLoader( MassSegmentationDataset(test_root_dir, combined_transform), batch_size=batch_size, num_workers=0 ) test(model, test_loader, device)
def main(): FLAGS = parser.parse_args() # Loading train and test data. # load_data function exists in file utils.py print("Loading dataset.") X_train, y_train = load_data(FLAGS.dataset_dir + '/train', FLAGS.img_size, FLAGS.augment_data) X_test, y_test = load_data(FLAGS.dataset_dir + '/test', FLAGS.img_size, FLAGS.augment_data) # Making sure that the data was loaded successfully. print("Train set image size : ", X_train.shape) print("Train set label size : ", y_train.shape) print("Test set image size : ", X_test.shape) print("Test set label size : ", y_test.shape) print("Dataset loaded successfully.") # Creating a unet object. # class UNet exists in file model.py unet = UNet(FLAGS.img_size) # Training the network, printing the loss value for every epoch # , and the accuracy on the test set after the training is complete train_loss_values, test_loss_values = unet.train(X_train, y_train, X_test, y_test, FLAGS.num_epochs, FLAGS.learning_rate, FLAGS.model_save_dir) # Plotting loss values on train and test set, and saving it as an image Loss.png # plot_loss exists in file utils.py plot_loss(train_loss_values, test_loss_values, 'Loss.png')
def main(args): train_dataloader, test_dataloader = dataloader.load_datasets( batch_size=args.batch_size, image_resize=args.image_resize, train_dataset_size=args.train_data_size, test_dataset_size=args.test_data_size, download=args.download_dataset ) model = UNet(out_channels=21) optimizer = optim.Adam(model.parameters(), lr=args.lr) ce_weight = utils.get_weight(train_dataloader.dataset) if len(ce_weight) < 21: criterion = nn.CrossEntropyLoss() else: criterion = nn.CrossEntropyLoss(utils.get_weight(train_dataloader.dataset)) print(f'Start training for {args.epochs} epochs') train(model=model, dataloader=train_dataloader, epochs=args.epochs, optimizer=optimizer, criterion=criterion, save_output_every=1, ) print(f'Training finished') print(f'Start evaluating with {len(test_dataloader.dataset)} images') eval(model, test_dataloader) print('All done')
def __init__(self, opt): self.opt = opt if opt.inference: self.testset = TestImageDataset(fdir=opt.impaths['test'], imsize=opt.imsize) else: self.trainset = ImageDataset(fdir=opt.impaths['train'], bdir=opt.impaths['btrain'], imsize=opt.imsize, mode='train', aug_prob=opt.aug_prob, prefetch=opt.prefetch) self.valset = ImageDataset(fdir=opt.impaths['val'], bdir=opt.impaths['bval'], imsize=opt.imsize, mode='val', aug_prob=opt.aug_prob, prefetch=opt.prefetch) self.model = UNet(n_channels=3, n_classes=1, bilinear=self.opt.use_bilinear) if opt.checkpoint: self.model.load_state_dict( torch.load('./state_dict/{:s}'.format(opt.checkpoint), map_location=self.opt.device)) print('checkpoint {:s} has been loaded'.format(opt.checkpoint)) if opt.multi_gpu == 'on': self.model = torch.nn.DataParallel(self.model) self.model = self.model.to(opt.device) self._print_args()
def main(): args = parser.parse_args() dataset = SyntheticCellDataset(arg.img_dir, arg.mask_dir) indices = torch.randperm(len(dataset)).tolist() sr = int(args.split_ratio * len(dataset)) train_set = torch.utils.data.Subset(dataset, indices[:-sr]) val_set = torch.utils.data.Subset(dataset, indices[-sr:]) train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True, pin_memory=True) val_loader = torch.utils.data.DataLoader(val_set, batch_size=args.batch_size, shuffle=False, pin_memory=True) device = torch.device("cpu" if not args.use_cuda else "cuda:0") model = UNet() model.to(device) dsc_loss = DiceLoss() optimizer = torch.optim.Adam(model.parameters(), args.lr) val_overall = 1000 for epoch in args.N_epoch: model, train_loss, optimizer = train(model, train_loader, device, optimizer) val_loss = validate(model, val_loader, device) if val_loss < val_overall: save_checkpoint(args.model_save_dir + '/epoch_'+str(epoch+1), model, train_loss, val_loss, epoch) val_overall = val_loss print('[{}/{}] train loss :{} val loss : {}'.format(epoch+1, num_epoch, train_loss, val_loss)) print('Training completed)
def lr_find(model: UNet, data_loader, optimizer: Optimizer, criterion, use_gpu, min_lr=0.0001, max_lr=0.1): # Save model and optimizer states to revert model_state = model.state_dict() optimizer_state = optimizer.state_dict() losses = [] lrs = [] scheduler = CyclicExpLR(optimizer, min_lr, max_lr, step_size_up=100, mode='triangular', cycle_momentum=True) model.train() for i, (data, target, class_ids) in enumerate(data_loader): data, target = data, target if use_gpu: data = data.cuda() target = target.cuda() optimizer.zero_grad() output_raw = model(data) # This step is specific for this project output = torch.zeros(output_raw.shape[0], 1, output_raw.shape[2], output_raw.shape[3]) if use_gpu: output = output.cuda() # This step is specific for this project for idx, (raw_o, class_id) in enumerate(zip(output_raw, class_ids)): output[idx] = raw_o[class_id - 1] loss = criterion(output, target) loss.backward() current_lr = optimizer.param_groups[0]['lr'] # Stop if lr stopped increasing if len(lrs) > 0 and current_lr < lrs[-1]: break lrs.append(current_lr) losses.append(loss.item()) optimizer.step() scheduler.step() # Plot in log scale plt.plot(lrs, losses) plt.xscale('log') plt.show() model.load_state_dict(model_state) optimizer.load_state_dict(optimizer_state)
def build_model(self): """Create a model""" self.model = UNet(self.in_dim, self.out_dim, self.num_filters) self.model = self.model.float() self.optimizer = torch.optim.Adam(self.model.parameters(), self.lr, [self.beta1, self.beta2], weight_decay=self.weight_decay) self.print_network(self.model, 'unet') self.model.to(self.device)
def main(): """ Training. """ global start_epoch, epoch, checkpoint # Initialize model or load checkpoint if checkpoint is None: model = UNet(in_channels, out_channels) # Initialize the optimizer optimizer = torch.optim.Adam(params=filter(lambda p: p.requires_grad, model.parameters()), lr=lr) else: checkpoint = torch.load(checkpoint) start_epoch = checkpoint['epoch'] + 1 model = checkpoint['model'] optimizer = checkpoint['optimizer'] # Move to default device model = model.to(device) criterion = nn.L1Loss().to(device) # Custom dataloaders train_loader = torch.utils.data.DataLoader(TripletDataset( train_folder, crop_size, scale), batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=True) test_loader = torch.utils.data.DataLoader(TripletDataset( test_folder, crop_size, scale), batch_size=batch_size, shuffle=True, num_workers=workers, pin_memory=True) # Total number of epochs to train for epochs = int(iterations // len(train_loader) + 1) # Epochs for epoch in range(start_epoch, epochs): # One epoch's training train(train_loader=train_loader, model=model, criterion=criterion, optimizer=optimizer, epoch=epoch, epochs=epochs) test(test_loader=test_loader, model=model, criterion=criterion) # Save checkpoint torch.save({ 'epoch': epoch, 'model': model, 'optimizer': optimizer }, f'checkpoints/checkpoint_unet_{epoch}.pth.tar')
def main(): parser = argparse.ArgumentParser() parser.add_argument("--action", type=str, default='train', help="train or test") args = parser.parse_args() config = load_config() # 使用tensorboard time_now = datetime.now().isoformat() if not os.path.exists(config.RUN_PATH): os.mkdir(config.RUN_PATH) writer = SummaryWriter(log_dir=config.RUN_PATH) # 随机数种子 torch.manual_seed(config.SEED) torch.cuda.manual_seed(config.SEED) np.random.seed(config.SEED) random.seed(config.SEED) # INIT GPU os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in config.GPU) if torch.cuda.is_available(): config.DEVICE = torch.device("cuda") print('\nGPU IS AVAILABLE') torch.backends.cudnn.benchmark = True else: config.DEVICE = torch.device("cpu") net = UNet(2).to(config.DEVICE) print(list(torchvision.models.resnet18(False).children())[7]) optimizer = optim.Adam(net.parameters(), betas=(0.5, 0.999), lr=config.LR) loss = nn.L1Loss() # 加载数据集 if args.action == 'train': train_dataset = LABDataset(config, config.TRAIN_PATH) len_train = len(train_dataset) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.BATCH_SIZE, shuffle=True) iter_per_epoch = len(train_loader) train_(config, train_loader, net, optimizer, loss, len_train, iter_per_epoch, writer) if args.action == "test": test_dataset = LABDataset(config, config.TEST_PATH) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False) test(config, test_loader, net, loss)
def __init__(self, cur_dir, suffix='.tif', cuda=True, testBatchSize=4, batchSize=4, nEpochs=200, lr=0.01, threads=4, seed=123, size=256, input_transform=True, target_transform=True): # super(TrainModel, self).__init__() self.data_dir = cur_dir + '/data/' self.suffix = suffix """ training parameters are set here """ self.colordim = 1 self.cuda = cuda if self.cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") self.testBatchSize = testBatchSize self.batchSize = batchSize self.nEpochs = nEpochs self.lr = lr self.threads = threads self.seed = seed self.size = size self.input_transform = input_transform self.target_transform = target_transform self.__check_dir = cur_dir + '/checkpoint' if not exists(self.__check_dir): os.mkdir(self.__check_dir) self.__epoch_dir = cur_dir + '/epoch' if not exists(self.__epoch_dir): os.mkdir(self.__epoch_dir) """ initialize the model """ if self.cuda: self.unet = UNet(self.colordim).cuda() self.criterion = nn.MSELoss().cuda() else: self.unet = UNet(self.colordim) self.criterion = nn.MSELoss() self.optimizer = optim.SGD(self.unet.parameters(), lr=self.lr, momentum=0.9, weight_decay=0.0001)
def main(): train_transform = A.Compose([ A.Resize(height=config.IMAGE_HEIGHT, width=config.IMAGE_WIDTH), A.Rotate(limit=35, p=1.0), A.HorizontalFlip(p=0.5), A.VerticalFlip(p=0.1), A.normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0], max_pixel_value=255.0), ToTensorV2, ]) val_transform = A.Compose([ A.Resize(height=config.IMAGE_HEIGHT, width=config.IMAGE_WIDTH), A.normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0], max_pixel_value=255.0), ToTensorV2, ]) model = UNet(in_channels=3, out_channels=1).to(config.DEVICE) loss_fn = nn.BCEWithLogitsLoss() optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE) train_loader, val_loader = get_loaders( config.TRAIN_IMAGE_DIR, config.TRAIN_MASK_DIR, config.VAL_IMG_DIR, config.VAL_MASK_DIR, config.BATCH_SIZE, train_transform, val_transform, ) if config.LOAD_MODEL: load_checkpoint(torch.load('my_checkpoint.pth.tar'), model) scaler = torch.cuda.amp.GradScaler() for epoch in range(config.NUM_EPOCHS): train_fn(train_loader, model, optimizer, loss_fn, scaler) # save model checkpoint = { 'state_dict': model.state_dict(), 'optimizer': optimizer.state_dict(), } save_checkpoint(checkpoint) # check acc check_accuracy(val_loader, model, device=config.DEVICE) # print some examples to a folder save_predictions_as_imgs(val_loader, model, folder='saved_images', device=config.DEVICE)
def main(): # width_in = 284 # height_in = 284 # width_out = 196 # height_out = 196 # PATH = './unet.pt' # x_train, y_train, x_val, y_val = get_dataset(width_in, height_in, width_out, height_out) # print(x_train.shape, y_train.shape, x_val.shape, y_val.shape) batch_size = 3 epochs = 1 epoch_lapse = 50 threshold = 0.5 learning_rate = 0.01 unet = UNet(in_channel=1, out_channel=2) if use_gpu: unet = unet.cuda() criterion = torch.nn.CrossEntropyLoss() optimizer = torch.optim.SGD(unet.parameters(), lr=0.01, momentum=0.99) if sys.argv[1] == 'train': train(unet, batch_size, epochs, epoch_lapse, threshold, learning_rate, criterion, optimizer, x_train, y_train, x_val, y_val, width_out, height_out) pass else: if use_gpu: unet.load_state_dict(torch.load(PATH)) else: unet.load_state_dict(torch.load(PATH, map_location='cpu')) print(unet.eval())
def create_net(self) -> nn.Module: if self.model_architecture == 'unet': self.net = UNet(n_channels=1, n_classes=1) elif self.model_architecture == 'mnet2': self.net = MobileNetV2_UNet() else: raise ValueError( f'model_architecture must be in ["unet", "mnet2"]. ' f'passed: {self.model_architecture}') self.net.to(device=self.device) return self.net
def __init__(self,hparams): super(Unetmodel,self).__init__() self.hparams = hparams self.hdf_path = self.hparams.hdf_path self.max_dist = self.hparams.max_dist self.grid_resolution = self.hparams.grid_resolution self.augment = self.hparams.augment self.net = UNet() self.metric = Ovl('Volume_overlap') self.loss = Dice_loss() self.lr = self.hparams.lr self.batch_size = self.hparams.batch_size
def __init__(self, device, n_channels=1, n_classes=1, load_model=None): self.device = device self.name = "unet" self.mask_type = torch.float32 if n_classes == 1 else torch.long self.net = UNet(n_channels=n_channels, n_classes=n_classes) if load_model is not None: self.net.load_state_dict(torch.load(load_model, map_location=device), strict=False) self.net = self.net.to(device) self.criterion = nn.BCELoss( ) if self.net.n_classes == 1 else nn.CrossEntropyLoss()
def main(mode): config = load_config() # 使用tensorboard time_now = datetime.now().isoformat() if not os.path.exists(config.RUN_PATH): os.mkdir(config.RUN_PATH) writer = SummaryWriter(log_dir=config.RUN_PATH) # 随机数种子 torch.manual_seed(config.SEED) torch.cuda.manual_seed(config.SEED) np.random.seed(config.SEED) random.seed(config.SEED) # INIT GPU os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in config.GPU) if torch.cuda.is_available(): config.DEVICE = torch.device("cuda") print('\nGPU IS AVAILABLE') torch.backends.cudnn.benchmark = True else: config.DEVICE = torch.device("cpu") net = UNet(1).to(config.DEVICE) optimizer = optim.Adam(net.parameters(), betas=(0.5, 0.999), lr=config.LR) #criterion = nn.CrossEntropyLoss() # 定义loss函数 criterion = nn.BCELoss() # 加载数据集 if mode == 1: train_dataset = MyDataset(config, config.TRAIN_PATH) len_train = len(train_dataset) train_loader = torch.utils.data.DataLoader( train_dataset, batch_size=config.BATCH_SIZE, shuffle=True) iter_per_epoch = len(train_loader) train_(config, train_loader, net, optimizer, criterion, len_train, iter_per_epoch, writer) if mode == 2: test_dataset = MyDataset(config, config.TEST_PATH) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1, shuffle=False) test(config, test_loader, net, criterion)
def build_model(self): """ Rough """ if self.cfg.task == 'cls': self.model = BinaryClassifier(num_classes=2) elif self.cfg.task == 'seg': self.model = UNet(num_classes=2) self.criterion = nn.CrossEntropyLoss() self.optim = torch.optim.Adam(self.model.parameters(), lr=self.cfg.lr, betas=(self.cfg.beta0, self.cfg.beta1)) if self.n_gpus > 1: print('### {} of gpus are used!!!'.format(self.n_gpus)) self.model = nn.DataParallel(self.model) self.model = self.model.to(self.device)
def get_score(): ckpt_dir = os.listdir( '/media/muyun99/DownloadResource/dataset/opends-Supervisely Person Dataset/checkpoints' ) Unet_dir = [ "25_noise", "50_noise", "75_noise", "25_noise_pro", "50_noise_pro", "75_noise_pro", "75_noise_pro_finetune" ] UNet_Pick_dir = ["75_noise_pro_QAM", "75_noise_pro_QAM_finetune"] UNet_Pick_cbam_dir = ["75_noise_pro_QAM_cbam_finetune"] for dir in ckpt_dir: print(dir) flag = "unet_pick" if dir in Unet_dir: net = UNet(n_classes=1, n_channels=3) flag = "unet" elif dir in UNet_Pick_dir: net = UNet_Pick(n_classes=1, n_channels=3) elif dir in UNet_Pick_cbam_dir: net = UNet_Pick_cbam(n_classes=1, n_channels=3) else: continue true_dir = os.path.join( '/media/muyun99/DownloadResource/dataset/opends-Supervisely Person Dataset/checkpoints', dir) try: best_ckpt_path = get_best_checkpoint(net, true_dir, flag) test_score = get_test_score(net, best_ckpt_path, flag) print(f'{dir} best test score is {test_score}') except Exception as ex: print(f'{dir} error') print(f'出现异常 {ex}') continue
def train(args): model = UNet(n_channels=125, n_classes=10).to(device) batch_size = args.batch_size num_epochs = args.num_epochs save_path = args.save_path criterion = L.CrossEntropyLoss2d() # 损失函数 optimizer = optim.Adam(model.parameters( )) # 梯度下降 # model.parameters():Returns an iterator over module parameters liver_dataset = LiverDataset("data", transform=x_transform, target_transform=y_transform) # 加载数据集 dataloader = DataLoader(liver_dataset, batch_size=batch_size, shuffle=True, num_workers=8) train_model(model, criterion, optimizer, dataloader, num_epochs, save_path) test(args)
def __init__(self,config,trainLoader,validLoader): self.config = config self.trainLoader = trainLoader self.validLoader = validLoader self.numTrain = len(self.trainLoader.dataset) self.numValid = len(self.validLoader.dataset) self.saveModelDir = str(self.config.save_model_dir)+"/" self.bestModel = config.bestModel self.useGpu = self.config.use_gpu self.net = UNet() if(self.config.resume == True): print("LOADING SAVED MODEL") self.loadCheckpoint() else: print("INTIALIZING NEW MODEL") self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") self.net = self.net.to(self.device) self.totalEpochs = config.epochs self.optimizer = optim.Adam(self.net.parameters(), lr=5e-4) self.loss = DiceLoss() self.num_params = sum([p.data.nelement() for p in self.net.parameters()]) self.trainPaitence = config.train_paitence if not self.config.resume: # self.freezeLayers(6) summary(self.net, input_size=(3,256,256)) print('[*] Number of model parameters: {:,}'.format(self.num_params)) self.writer = SummaryWriter(self.config.tensorboard_path+"/")
def __init__(self, n_channels, n_classes): super(UNet_Pick_cbam, self).__init__() self.n_channels = n_channels self.n_classes = n_classes self.Unet = UNet(n_classes=n_classes, n_channels=n_channels) self.QAM = implement.QAM_cbam() self.OCM = implement.OCM()
def main(conf): device = "cuda:0" if torch.cuda.is_available() else 'cpu' beta_schedule = "linear" beta_start = 1e-4 beta_end = 2e-2 n_timestep = 1000 conf.distributed = dist.get_world_size() > 1 transform = transforms.Compose( [ transforms.RandomHorizontalFlip(), transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True), ] ) train_set = MultiResolutionDataset( conf.dataset.path, transform, conf.dataset.resolution ) train_sampler = dist.data_sampler( train_set, shuffle=True, distributed=conf.distributed ) train_loader = conf.training.dataloader.make(train_set, sampler=train_sampler) model = UNet( conf.model.in_channel, conf.model.channel, channel_multiplier=conf.model.channel_multiplier, n_res_blocks=conf.model.n_res_blocks, attn_strides=conf.model.attn_strides, dropout=conf.model.dropout, fold=conf.model.fold, ) model = model.to(device) ema = UNet( conf.model.in_channel, conf.model.channel, channel_multiplier=conf.model.channel_multiplier, n_res_blocks=conf.model.n_res_blocks, attn_strides=conf.model.attn_strides, dropout=conf.model.dropout, fold=conf.model.fold, ) ema = ema.to(device) if conf.distributed: model = nn.parallel.DistributedDataParallel( model, device_ids=[dist.get_local_rank()], output_device=dist.get_local_rank(), ) optimizer = conf.training.optimizer.make(model.parameters()) scheduler = conf.training.scheduler.make(optimizer) betas = make_beta_schedule(beta_schedule, beta_start, beta_end, n_timestep) diffusion = GaussianDiffusion(betas).to(device) train(conf, train_loader, model, ema, diffusion, optimizer, scheduler, device)
def load_model(data, model_path, cuda=True): if cuda and not torch.cuda.is_available(): raise Exception("No GPU found, please run without --cuda") unet = UNet() if cuda: unet = unet.cuda() if not cuda: unet.load_state_dict( torch.load(model_path, map_location=lambda storage, loc: storage)) else: unet.load_state_dict(torch.load(model_path)) if cuda: data = Variable(data.cuda()) else: data = Variable(data) data = torch.unsqueeze(data, 0) output = unet(data) if cuda: output = output.cuda() return output
def main(): args = parser.parse_args() device = torch.device("cpu" if not args.use_cuda else "cuda:0") model = UNet() load_checkpoint(args.weight, model, device='cpu') model.to(device) img = Image.open(img_file) resize = transforms.Resize(size=(576, 576)) im_r = TF.to_tensor(resize(img)) im_r = im_r.unsqueeze(0) with torch.no_grad(): pred = model(im_r.to(device)) pred_mask = pred.detach().cpu().numpy().squeeze()
def main(tocsv=False, save=False, mask=False, valid_train=False, toiou=False): model_name = config['param']['model'] resize = not config['valid'].getboolean('pred_orig_size') if model_name == 'unet_vgg16': model = UNetVgg16(3, 1, fixed_vgg=True) elif model_name == 'dcan': model = DCAN(3, 1) elif model_name == 'caunet': model = CAUNet() elif model_name == 'camunet': model = CAMUNet() else: model = UNet() if torch.cuda.is_available(): model = model.cuda() # model = torch.nn.DataParallel(model).cuda() # Sets the model in evaluation mode. model.eval() epoch = load_ckpt(model) if epoch == 0: print("Aborted: checkpoint not found!") return # prepare dataset compose = Compose(augment=False, resize=resize) data_dir = 'data/stage1_train' if valid_train else 'data/stage1_test' dataset = KaggleDataset(data_dir, transform=compose) iter = predict(model, dataset, compose, resize) if tocsv: with open('result.csv', 'w') as csvfile: writer = csv.writer(csvfile) writer.writerow(['ImageId', 'EncodedPixels']) for uid, _, y, y_c, y_m, _, _, _, _ in iter: for rle in prob_to_rles(y, y_c, y_m): writer.writerow([uid, ' '.join([str(i) for i in rle])]) elif toiou and valid_train: with open('iou.csv', 'w') as csvfile: writer = csv.writer(csvfile) writer.writerow(['ImageId', 'IoU']) for uid, _, y, y_c, y_m, gt, _, _, _ in tqdm(iter): iou = get_iou(y, y_c, y_m, gt) writer.writerow([uid, iou]) else: for uid, x, y, y_c, y_m, gt, gt_s, gt_c, gt_m in tqdm(iter): if valid_train: show_groundtruth(uid, x, y, y_c, y_m, gt, gt_s, gt_c, gt_m, save) elif mask: save_mask(uid, y, y_c, y_m) else: show(uid, x, y, y_c, y_m, save)