\n\nIn this tutorial, we use InversionNet ([Wu and Lin, 2019](https://ieeexplore.ieee.org/abstract/document/8918045)):\n\nThe network structure (see *network.py*)\n\n
\n\n\nThe network is mainly composed of 2D convolution layers and transposed 2D convolution layers.\n\n\n\n*2D convolution with no padding, stride of 2 and kernel of 3. Image source: https://www.analyticsvidhya.com/blog/2022/03/basics-of-cnn-in-deep-learning/\n\n\n\n\n*Transposed 2D convolution with no padding, stride of 2 and kernel of 3.Image source: https://towardsdatascience.com/types-of-convolutions-in-deep-learning-717013397f4d\n","metadata":{}},{"cell_type":"code","source":"%%writefile dataset.py\nimport os\nimport numpy as np\nfrom torch.utils.data import Dataset\nfrom torchvision.transforms import Compose\nimport transforms as T\n\nclass FWIDataset(Dataset):\n def __init__(self, anno, preload=True, sample_ratio=1, file_size=500,\n transform_data=None, transform_label=None):\n if not os.path.exists(anno):\n print(f'Annotation file {anno} does not exists')\n self.preload = preload\n self.sample_ratio = sample_ratio\n self.file_size = file_size\n self.transform_data = transform_data\n self.transform_label = transform_label\n with open(anno, 'r') as f:\n self.batches = f.readlines()\n if preload: \n self.dataset_list, self.data_list, self.label_list = [], [], []\n for batch in self.batches: \n dataset, data, label = self.load_every(batch)\n self.dataset_list.append(dataset)\n self.data_list.append(data)\n if label is not None:\n self.label_list.append(label)\n\n # Load from one line\n def load_every(self, batch):\n batch = batch.strip().split('\\t')\n dataset = batch[0]\n data_path = batch[1]\n data = np.load(data_path)[:, :, ::self.sample_ratio, :]\n data = data.astype('float32')\n if len(batch) > 2:\n label_path = batch[2] \n label = np.load(label_path)\n label = label.astype('float32')\n else:\n label = None\n \n return dataset, data, label\n \n def __getitem__(self, idx):\n batch_idx, sample_idx = idx // self.file_size, idx % self.file_size\n if self.preload:\n dataset = self.dataset_list[batch_idx][sample_idx]\n data = self.data_list[batch_idx][sample_idx]\n label = self.label_list[batch_idx][sample_idx] if len(self.label_list) != 0 else None\n else:\n dataset, data, label = self.load_every(self.batches[batch_idx])\n data = data[sample_idx]\n label = label[sample_idx] if label is not None else None\n if self.transform_data:\n data = self.transform_data[dataset](data)\n if self.transform_label and label is not None:\n label = self.transform_label[dataset](label)\n return data, label if label is not None else np.array([])\n \n def __len__(self):\n return len(self.batches) * self.file_size","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T19:13:27.986819Z","iopub.execute_input":"2025-04-12T19:13:27.987125Z","iopub.status.idle":"2025-04-12T19:13:27.993948Z","shell.execute_reply.started":"2025-04-12T19:13:27.987100Z","shell.execute_reply":"2025-04-12T19:13:27.993284Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%writefile train.py\n\n\n\nimport os\nimport sys\nimport time\nimport datetime\nimport json\n\nimport torch\nfrom torch import nn\nfrom torch.utils.data import RandomSampler, DataLoader\nfrom torch.utils.data.dataloader import default_collate\nfrom torch.utils.data.distributed import DistributedSampler\nfrom torch.nn.parallel import DistributedDataParallel\nfrom torch.utils.tensorboard import SummaryWriter\nimport torchvision\nfrom torchvision.transforms import Compose\n\nimport utils\nimport network\nfrom dataset import FWIDataset\nfrom scheduler import WarmupMultiStepLR\nimport transforms as T\n\nstep = 0\n\ndef train_one_epoch(model, criterion, optimizer, lr_scheduler, \n dataloader, device, epoch, print_freq, writer):\n global step\n model.train()\n\n # Logger setup\n metric_logger = utils.MetricLogger(delimiter=' ')\n metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value}'))\n metric_logger.add_meter('samples/s', utils.SmoothedValue(window_size=10, fmt='{value:.3f}'))\n header = 'Epoch: [{}]'.format(epoch)\n\n for data, label in metric_logger.log_every(dataloader, print_freq, header):\n start_time = time.time()\n optimizer.zero_grad()\n data, label = data.to(device), label.to(device)\n output = model(data)\n loss, loss_g1v, loss_g2v = criterion(output, label)\n loss.backward()\n optimizer.step()\n\n loss_val = loss.item()\n loss_g1v_val = loss_g1v.item()\n loss_g2v_val = loss_g2v.item()\n batch_size = data.shape[0]\n metric_logger.update(loss=loss_val, loss_g1v=loss_g1v_val, \n loss_g2v=loss_g2v_val, lr=optimizer.param_groups[0]['lr'])\n metric_logger.meters['samples/s'].update(batch_size / (time.time() - start_time))\n if writer:\n writer.add_scalar('loss', loss_val, step)\n writer.add_scalar('loss_g1v', loss_g1v_val, step)\n writer.add_scalar('loss_g2v', loss_g2v_val, step)\n step += 1\n lr_scheduler.step()\n\n\ndef evaluate(model, criterion, dataloader, device, writer):\n model.eval()\n metric_logger = utils.MetricLogger(delimiter=' ')\n header = 'Test:'\n with torch.no_grad():\n for data, label in metric_logger.log_every(dataloader, 20, header):\n data = data.to(device, non_blocking=True)\n label = label.to(device, non_blocking=True)\n output = model(data)\n loss, loss_g1v, loss_g2v = criterion(output, label)\n metric_logger.update(loss=loss.item(), \n loss_g1v=loss_g1v.item(), \n loss_g2v=loss_g2v.item())\n\n # Gather the stats from all processes\n metric_logger.synchronize_between_processes()\n print(' * Loss {loss.global_avg:.8f}\\n'.format(loss=metric_logger.loss))\n if writer:\n writer.add_scalar('loss', metric_logger.loss.global_avg, step)\n writer.add_scalar('loss_g1v', metric_logger.loss_g1v.global_avg, step)\n writer.add_scalar('loss_g2v', metric_logger.loss_g2v.global_avg, step)\n return metric_logger.loss.global_avg\n\n\ndef main(args):\n global step\n\n print(args)\n print('torch version: ', torch.__version__)\n print('torchvision version: ', torchvision.__version__)\n\n utils.mkdir(args.output_path) # create folder to store checkpoints\n utils.init_distributed_mode(args) # distributed mode initialization\n \n # Set up tensorboard summary writer\n train_writer, val_writer = None, None\n if args.tensorboard:\n utils.mkdir(args.log_path) # create folder to store tensorboard logs\n if not args.distributed or (args.rank == 0) and (args.local_rank == 0):\n train_writer = SummaryWriter(os.path.join(args.output_path, 'logs', 'train'))\n val_writer = SummaryWriter(os.path.join(args.output_path, 'logs', 'val'))\n \n\n device = torch.device(args.device)\n torch.backends.cudnn.benchmark = True\n transform_data, transform_label = {}, {}\n \n with open('dataset_config.json') as f:\n try:\n ctx_dict = json.load(f)\n for dataset, ctx in ctx_dict.items():\n # Normalize data and label to [-1, 1]\n transform_data[dataset] = Compose([\n T.LogTransform(k=args.k),\n T.MinMaxNormalize(T.log_transform(ctx['data_min'], k=args.k), T.log_transform(ctx['data_max'], k=args.k))\n ])\n transform_label[dataset] = Compose([\n T.MinMaxNormalize(ctx['label_min'], ctx['label_max'])\n ])\n except KeyError:\n print('Unsupported dataset.')\n sys.exit()\n\n # Create dataset and dataloader\n print('Loading data')\n print('Loading training data')\n file_size = 500\n if args.file_size is not None:\n file_size = args.file_size\n if args.train_anno[-3:] == 'txt':\n dataset_train = FWIDataset(\n args.train_anno,\n preload=False,\n sample_ratio=args.sample_temporal,\n file_size=file_size,\n transform_data=transform_data,\n transform_label=transform_label\n )\n else:\n dataset_train = torch.load(args.train_anno)\n\n print('Loading validation data')\n if args.val_anno[-3:] == 'txt':\n dataset_valid = FWIDataset(\n args.val_anno,\n preload=False,\n sample_ratio=args.sample_temporal,\n file_size=file_size,\n transform_data=transform_data,\n transform_label=transform_label\n )\n else:\n dataset_valid = torch.load(args.val_anno)\n\n print('Creating data loaders')\n if args.distributed:\n train_sampler = DistributedSampler(dataset_train, shuffle=True)\n valid_sampler = DistributedSampler(dataset_valid, shuffle=True)\n else:\n train_sampler = RandomSampler(dataset_train)\n valid_sampler = RandomSampler(dataset_valid)\n\n dataloader_train = DataLoader(\n dataset_train, batch_size=args.batch_size,\n sampler=train_sampler, num_workers=args.workers,\n pin_memory=True, drop_last=True, collate_fn=default_collate)\n\n dataloader_valid = DataLoader(\n dataset_valid, batch_size=args.batch_size,\n sampler=valid_sampler, num_workers=args.workers,\n pin_memory=True, collate_fn=default_collate)\n\n print('Creating model')\n if args.model not in network.model_dict:\n print('Unsupported model.')\n sys.exit()\n model = network.model_dict[args.model](upsample_mode=args.up_mode, \n sample_spatial=args.sample_spatial, sample_temporal=args.sample_temporal).to(device)\n\n if args.distributed and args.sync_bn:\n model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)\n\n # Define loss function\n l1loss = nn.L1Loss()\n l2loss = nn.MSELoss()\n def criterion(pred, gt):\n loss_g1v = l1loss(pred, gt)\n loss_g2v = l2loss(pred, gt)\n loss = args.lambda_g1v * loss_g1v + args.lambda_g2v * loss_g2v\n return loss, loss_g1v, loss_g2v\n\n # Scale lr according to effective batch size\n lr = args.lr * args.world_size\n optimizer = torch.optim.AdamW(model.parameters(), lr=lr, betas=(0.9, 0.999), weight_decay=args.weight_decay)\n\n # Convert scheduler to be per iteration instead of per epoch\n warmup_iters = args.lr_warmup_epochs * len(dataloader_train)\n lr_milestones = [len(dataloader_train) * m for m in args.lr_milestones]\n lr_scheduler = WarmupMultiStepLR(\n optimizer, milestones=lr_milestones, gamma=args.lr_gamma,\n warmup_iters=warmup_iters, warmup_factor=1e-5)\n\n model_without_ddp = model\n if args.distributed:\n model = DistributedDataParallel(model, device_ids=[args.local_rank])\n model_without_ddp = model.module\n\n if args.resume:\n checkpoint = torch.load(args.resume, map_location='cpu')\n model_without_ddp.load_state_dict(network.replace_legacy(checkpoint['model']))\n optimizer.load_state_dict(checkpoint['optimizer'])\n lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])\n args.start_epoch = checkpoint['epoch'] + 1\n step = checkpoint['step']\n lr_scheduler.milestones=lr_milestones\n\n print('Start training')\n start_time = time.time()\n best_loss = 10\n chp=1 \n for epoch in range(args.start_epoch, args.epochs):\n if args.distributed:\n train_sampler.set_epoch(epoch)\n train_one_epoch(model, criterion, optimizer, lr_scheduler, dataloader_train,\n device, epoch, args.print_freq, train_writer)\n \n loss = evaluate(model, criterion, dataloader_valid, device, val_writer)\n \n checkpoint = {\n 'model': model_without_ddp.state_dict(),\n 'optimizer': optimizer.state_dict(),\n 'lr_scheduler': lr_scheduler.state_dict(),\n 'epoch': epoch,\n 'step': step,\n 'args': args}\n # Save checkpoint per epoch\n if loss < best_loss:\n utils.save_on_master(\n checkpoint,\n os.path.join(args.output_path, 'checkpoint.pth'))\n print('saving checkpoint at epoch: ', epoch)\n chp = epoch\n best_loss = loss\n # Save checkpoint every epoch block\n print('current best loss: ', best_loss)\n print('current best epoch: ', chp)\n if args.output_path and (epoch + 1) % args.epoch_block == 0:\n utils.save_on_master(\n checkpoint,\n os.path.join(args.output_path, 'model_{}.pth'.format(epoch + 1)))\n\n total_time = time.time() - start_time\n total_time_str = str(datetime.timedelta(seconds=int(total_time)))\n print('Training time {}'.format(total_time_str))\n\n\ndef parse_args():\n import argparse\n parser = argparse.ArgumentParser(description='FCN Training')\n parser.add_argument('-d', '--device', default='cuda', help='device')\n #parser.add_argument('-ds', '--dataset', default='flatfault-b', type=str, help='dataset name')\n parser.add_argument('-fs', '--file-size', default=None, type=int, help='number of samples in each npy file')\n\n # Path related\n parser.add_argument('-ap', '--anno-path', default='split_files', help='annotation files location')\n parser.add_argument('-t', '--train-anno', default='flatfault_b_train_invnet.txt', help='name of train anno')\n parser.add_argument('-v', '--val-anno', default='flatfault_b_val_invnet.txt', help='name of val anno')\n parser.add_argument('-o', '--output-path', default='Invnet_models', help='path to parent folder to save checkpoints')\n parser.add_argument('-l', '--log-path', default='Invnet_models', help='path to parent folder to save logs')\n parser.add_argument('-n', '--save-name', default='fcn_l1loss_ffb', help='folder name for this experiment')\n parser.add_argument('-s', '--suffix', type=str, default=None, help='subfolder name for this run')\n\n # Model related\n parser.add_argument('-m', '--model', type=str, help='inverse model name')\n parser.add_argument('-um', '--up-mode', default=None, help='upsampling layer mode such as \"nearest\", \"bicubic\", etc.')\n parser.add_argument('-ss', '--sample-spatial', type=float, default=1.0, help='spatial sampling ratio')\n parser.add_argument('-st', '--sample-temporal', type=int, default=1, help='temporal sampling ratio')\n # Training related\n parser.add_argument('-b', '--batch-size', default=256, type=int)\n parser.add_argument('--lr', default=0.0001, type=float, help='initial learning rate')\n parser.add_argument('-lm', '--lr-milestones', nargs='+', default=[], type=int, help='decrease lr on milestones')\n parser.add_argument('--momentum', default=0.9, type=float, help='momentum')\n parser.add_argument('--weight-decay', default=1e-4 , type=float, help='weight decay (default: 1e-4)')\n parser.add_argument('--lr-gamma', default=0.1, type=float, help='decrease lr by a factor of lr-gamma')\n parser.add_argument('--lr-warmup-epochs', default=0, type=int, help='number of warmup epochs') \n parser.add_argument('-eb', '--epoch_block', type=int, default=40, help='epochs in a saved block')\n parser.add_argument('-nb', '--num_block', type=int, default=3, help='number of saved block')\n parser.add_argument('-j', '--workers', default=4, type=int, help='number of data loading workers (default: 16)')\n parser.add_argument('--k', default=1, type=float, help='k in log transformation')\n parser.add_argument('--print-freq', default=50, type=int, help='print frequency')\n parser.add_argument('-r', '--resume', default=None, help='resume from checkpoint')\n parser.add_argument('--start-epoch', default=0, type=int, help='start epoch')\n\n # Loss related\n parser.add_argument('-g1v', '--lambda_g1v', type=float, default=1.0)\n parser.add_argument('-g2v', '--lambda_g2v', type=float, default=1.0)\n \n # Distributed training related\n parser.add_argument('--sync-bn', action='store_true', help='Use sync batch norm')\n parser.add_argument('--world-size', default=2, type=int, help='number of distributed processes')\n parser.add_argument('--dist-url', default='env://', help='url used to set up distributed training')\n\n # Tensorboard related\n parser.add_argument('--tensorboard', action='store_true', help='Use tensorboard for logging.')\n\n args = parser.parse_args()\n\n args.output_path = os.path.join(args.output_path, args.save_name, args.suffix or '')\n args.log_path = os.path.join(args.log_path, args.save_name, args.suffix or '')\n args.train_anno = os.path.join(args.anno_path, args.train_anno)\n args.val_anno = os.path.join(args.anno_path, args.val_anno)\n \n args.epochs = args.epoch_block * args.num_block\n\n if args.resume:\n args.resume = os.path.join(args.output_path, args.resume)\n\n return args\n\n\nif __name__ == '__main__':\n args = parse_args()\n main(args)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T19:13:29.531689Z","iopub.execute_input":"2025-04-12T19:13:29.532292Z","iopub.status.idle":"2025-04-12T19:13:29.542065Z","shell.execute_reply.started":"2025-04-12T19:13:29.532268Z","shell.execute_reply":"2025-04-12T19:13:29.541389Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Train InversionNet","metadata":{}},{"cell_type":"code","source":"!python train.py -n openfwi_cnn_1 -m InversionNet -g1v 1 -g2v 0 --tensorboard -ap ./inputs -t train.txt -v val.txt --lr 0.0001 -b 256 -eb 1 -nb 1 -fs 100","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T19:15:34.865329Z","iopub.execute_input":"2025-04-12T19:15:34.866186Z","iopub.status.idle":"2025-04-12T19:36:26.221450Z","shell.execute_reply.started":"2025-04-12T19:15:34.866156Z","shell.execute_reply":"2025-04-12T19:36:26.220685Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! ls -lsh Invnet_models/openfwi_cnn_1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T19:42:26.416155Z","iopub.execute_input":"2025-04-12T19:42:26.416964Z","iopub.status.idle":"2025-04-12T19:42:26.600820Z","shell.execute_reply.started":"2025-04-12T19:42:26.416925Z","shell.execute_reply":"2025-04-12T19:42:26.600131Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}
![](\"https://openfwi-lanl.github.io/assets/img/InversionNet.png\")
![](\"https://editor.analyticsvidhya.com/uploads/33383str.jpg\")
![](\"https://miro.medium.com/max/790/1*Lpn4nag_KRMfGkx1k6bV-g.gif\")