def create_datasets(args): train_mask = create_mask_for_mask_type(args.mask_type, args.center_fractions, args.accelerations) dev_mask = create_mask_for_mask_type(args.mask_type, args.center_fractions, args.accelerations) train_data = SliceData( root=args.data_path / f'{args.challenge}_train', transform=DataTransform(train_mask, args.resolution, args.challenge), sample_rate=args.sample_rate, challenge=args.challenge ) dev_data = SliceData( root=args.data_path / f'{args.challenge}_val', transform=DataTransform(dev_mask, args.resolution, args.challenge, use_seed=True), sample_rate=args.sample_rate, challenge=args.challenge, ) return dev_data, train_data
def train_data_transform(self): mask = create_mask_for_mask_type(self.hparams.mask_type, self.hparams.center_fractions, self.hparams.accelerations) return DataTransform(self.hparams.resolution, self.hparams.challenge, mask, use_seed=False)
def create_data_loader(args): dev_mask = create_mask_for_mask_type(args.mask_type, args.center_fractions, args.accelerations) data = SliceData( root=args.data_path / f'{args.challenge}_val', transform=DataTransform(dev_mask), challenge=args.challenge, sample_rate=args.sample_rate ) return data
def __init__(self, reg_network=None, hparams=None): super(NeumannNetwork, self).__init__() self.hparams = hparams self.device = "cuda" if hparams.gpus == 0: self.device = "cpu" self.mask_func = create_mask_for_mask_type( self.hparams.mask_type, self.hparams.center_fractions, self.hparams.accelerations) self.reg_network = reg_network self.n_blocks = hparams.n_blocks self.eta = nn.Parameter(torch.Tensor([0.1]), requires_grad=True) self.preconditioned = False
def create_data_loaders(args): mask_func = None if args.mask_kspace: mask_func = create_mask_for_mask_type(args.mask_type, args.center_fractions, args.accelerations) data = SliceData( root=args.data_path / f'{args.challenge}_{args.data_split}', transform=DataTransform(args.resolution, args.challenge, mask_func), sample_rate=1., challenge=args.challenge ) data_loader = DataLoader( dataset=data, batch_size=args.batch_size, num_workers=4, pin_memory=True, ) return data_loader
def validation_step(self, batch, batch_idx): # print(f"Validation step, batch_idx:{batch_idx}") image, target, kspace, mean, std, fname, slice = batch output = self.forward(kspace) mask = create_mask_for_mask_type(self.hparams.mask_type, self.hparams.center_fractions, self.hparams.accelerations) _, undersampled_img = forward_adjoint_helper(self.device, self.hparams, mask, kspace, None) mean = mean.unsqueeze(1).unsqueeze(2) std = std.unsqueeze(1).unsqueeze(2) return { "fname": fname, "slice": slice, "output": (output * std + mean).cpu().numpy(), "target": (target * std + mean).cpu().numpy(), "undersampled_img": (undersampled_img * std + mean).cpu().numpy(), "val_loss": F.l1_loss(output, target), }
def val_data_transform(self): mask = create_mask_for_mask_type(self.hparams.mask_type, self.hparams.center_fractions, self.hparams.accelerations) return DataTransform(mask, self.hparams.resolution)
def train_epoch(args, epoch, model, data_loader, optimizer, writer): model.train() avg_loss = 0. start_epoch = start_iter = time.perf_counter() global_step = epoch * len(data_loader) if args.fn_train: #loss_f = torch.nn.MSELoss(reduction='none') loss_f = torch.nn.L1Loss(reduction='none') mask_f = create_mask_for_mask_type('random',[0.08, 0.04], [4, 8]) n_pixel_range = (10, 100) n_pixel_range = (args.min_n_pixel, args.max_n_pixel) for iter, data in enumerate(data_loader): if args.bbox_root: (input, target, mean, std, norm), seg = data else: if not args.fn_train: input, target, mean, std, norm = data else: input, target, mean, std, norm, fn_image = data input = input.unsqueeze(1).to(args.device) target = target.to(args.device) output = model(input).squeeze(1) loss = F.l1_loss(output, target) if args.bbox_root: writer.add_scalar('L1_Loss', loss.item(), global_step + iter) bbox_loss = [] for j in range(11): seg_mask = seg[:, :, :, j] if seg_mask.sum() > 0: seg_mask = seg_mask.to(args.device) bbox_output = output * seg_mask bbox_target = target * seg_mask bbox_loss.append(nmse(bbox_target, bbox_output)) if len(bbox_loss)>0: bbox_loss = 2 * torch.stack(bbox_loss).mean() #print(loss.item(), bbox_loss.item()) writer.add_scalar('BBOX_Loss', bbox_loss.item(), global_step + iter) loss += bbox_loss optimizer.zero_grad() loss.backward() optimizer.step() #information_loss_list, xs_list, ys_list = [], [], [] if args.fn_train: #fn_attack_train(model, target, optimizer) #run_finite_diff(model, target, iterations=2, train=True, optimizer=optimizer) get_attack_loss_new(model, fn_image, loss_f=loss_f, xs=np.random.randint(low=from_boarder, high=320-from_boarder, size=(fn_image.size(0),)), ys=np.random.randint(low=from_boarder, high=320-from_boarder, size=(fn_image.size(0),)), shape=(320, 320), n_pixel_range=n_pixel_range, train=True, optimizer=optimizer) avg_loss = 0.99 * avg_loss + 0.01 * loss.item() if iter > 0 else loss.item() writer.add_scalar('TrainLoss', loss.item(), global_step + iter) if iter % args.report_interval == 0: logging.info( f'Epoch = [{epoch:3d}/{args.num_epochs:3d}] ' f'Iter = [{iter:4d}/{len(data_loader):4d}] ' f'Loss = {loss.item():.4g} Avg Loss = {avg_loss:.4g} ' f'Time = {time.perf_counter() - start_iter:.4f}s', ) start_iter = time.perf_counter() return avg_loss, time.perf_counter() - start_epoch
parser.add_argument('--bbox_root', type=str, default=None) parser.add_argument('--fnaf_eval', type=pathlib.Path, default=None, help='Path where fnaf eval results should be saved') parser.add_argument('--fnaf_eval_control', action='store_true') return parser if __name__ == '__main__': args = create_arg_parser().parse_args() import os os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID" # The GPU id to use, usually either "0" or "1" os.environ["CUDA_VISIBLE_DEVICES"]=args.gpu mask_f = create_mask_for_mask_type(args.mask_type, args.center_fractions, args.accelerations) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) main(args)