cmap=args.cmap) if epoch % args.log_freq == 0: logger['psnr_test'].append(psnr) logger['rmse_test'].append(rmse) print("Epoch {}: test PSNR: {:.6f}, RMSE: {:.6f}".format( epoch, psnr, rmse)) tic2 = time.time() print("Finished training {} epochs using {} seconds".format( args.epochs, tic2 - tic)) x_axis = np.arange(args.log_freq, args.epochs + args.log_freq, args.log_freq) # plot the rmse, r2-score curve and save them in txt save_stats(args.train_dir, logger, x_axis, 'psnr_train', 'psnr_test', 'rmse_train', 'rmse_test') args.training_time = tic2 - tic args.n_params, args.n_layers = module_size(model) with open(args.run_dir + "/args.txt", 'w') as args_file: json.dump(vars(args), args_file, indent=4) except KeyboardInterrupt: print('Keyboard Interrupt captured...Saving models & training logs') tic2 = time.time() torch.save(model.state_dict(), args.ckpt_dir + "/model_epoch{}.pth".format(epoch)) x_axis = np.arange(args.log_freq, args.epochs + args.log_freq, args.log_freq) # plot the rmse, r2-score curve and save them in txt save_stats(args.train_dir, logger, x_axis, 'psnr_train', 'psnr_test',
print( f'Epoch {epoch}: training loss: {loss_train:.6f} ' \ f'relative-ux: {relative_ux: .5f}, relative-uy: {relative_uy: .5f}') logf = open(args.train_dir + "/running_log.txt", 'a') logf.write( f'Epoch {epoch}: training loss: {loss_train:.6f}, ' \ f'relative-ux: {relative_ux: .5f}, relative-uy: {relative_uy: .5f} \n') logf.close() if epoch % args.log_freq == 0: logger['loss_train'].append(loss_train) logger['ux_l2loss'].append(relative_ux) logger['uy_l2loss'].append(relative_uy) if epoch % args.ckpt_freq == 0: torch.save(model, args.ckpt_dir + "/model_epoch{}.pth".format(epoch)) sampledir = args.ckpt_dir + "/model{}".format(epoch) mkdirs((sampledir)) tic2 = time.time() print(f'Finished training {args.epochs} epochs with {args.ntrain} data ' \ f'using {(tic2 - tic) / 60:.2f} mins') metrics = ['loss_train', 'ux_l2loss', 'uy_l2loss'] save_stats(args.train_dir, logger, *metrics) args.training_time = tic2 - tic args.n_params, args.n_layers = model.model_size with open(args.run_dir + "/args.txt", 'w') as args_file: json.dump(vars(args), args_file, indent=4)
logger['r2_test'].append(r2_test) logger['mnlp_test'].append(mnlp_test) print("epoch {}, testing r2: {:.6f}, test mnlp: {}".format( epoch, r2_test, mnlp_test)) if __name__ == '__main__': print( 'Start training.........................................................' ) tic = time() for epoch in range(1, args.epochs + 1): svgd.train(epoch, logger) with torch.no_grad(): test(epoch, logger) training_time = time() - tic print( 'Finished training:\n{} epochs\n{} data\n{} samples (SVGD)\n{} seconds' .format(args.epochs, args.ntrain, args.n_samples, training_time)) # save training results x_axis = np.arange(args.log_freq, args.epochs + args.log_freq, args.log_freq) # plot the rmse, r2-score curve and save them in txt save_stats(args.train_dir, logger, x_axis) args.training_time = training_time args.n_params, args.n_layers = dense_ed._num_parameters_convlayers() with open(args.run_dir + "/args.txt", 'w') as args_file: json.dump(vars(args), args_file, indent=4)
def main(): parser = argparse.ArgumentParser(description='CNN to solve PDE') parser.add_argument('--exp-dir', type=str, default='./experiments/solver', help='color map') parser.add_argument('--nonlinear', action='store_true', default=False, help='set True for nonlinear PDE') # data parser.add_argument('--data-dir', type=str, default="./datasets", help='directory to dataset') parser.add_argument('--data', type=str, default='grf', choices=['grf', 'channelized', 'warped_grf'], help='data type') parser.add_argument('--kle', type=int, default=512, help='# kle terms') parser.add_argument('--imsize', type=int, default=64, help='image size') parser.add_argument('--idx', type=int, default=8, help='idx of input, please use 0 ~ 999') parser.add_argument('--alpha1', type=float, default=1.0, help='coefficient for the squared term') parser.add_argument('--alpha2', type=float, default=1.0, help='coefficient for the cubic term') # latent size: (nz, sz, sz) parser.add_argument('--nz', type=int, default=1, help='# feature maps of latent z') # parser.add_argument('--sz', type=int, default=16, help='feature map size of latent z') parser.add_argument('--blocks', type=list, default=[8, 6], help='# layers in each dense block of the decoder') parser.add_argument('--weight-bound', type=float, default=10, help='weight for boundary condition loss') parser.add_argument('--lr', type=float, default=0.5, help='learning rate') parser.add_argument('--epochs', type=int, default=500, help='# epochs to train') parser.add_argument('--test-freq', type=int, default=50, help='every # epoch to test') parser.add_argument('--ckpt-freq', type=int, default=250, help='every # epoch to save model') parser.add_argument('--cmap', type=str, default='jet', help='color map') parser.add_argument( '--same-scale', action='store_true', help='true for setting noise to be same scale as output') parser.add_argument('--animate', action='store_true', help='true to plot animate figures') parser.add_argument('--cuda', type=int, default=1, help='cuda number') parser.add_argument('-v', '--verbose', action='store_true', help='True for versbose output') args = parser.parse_args() pprint(vars(args)) device = torch.device( f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu") dataset = f'{args.data}_kle{args.kle}' if args.data == 'grf' else args.data hyparams = f'{dataset}_idx{args.idx}_dz{args.nz}_blocks{args.blocks}_'\ f'lr{args.lr}_wb{args.weight_bound}_epochs{args.epochs}' if args.nonlinear: from utils.fenics import solve_nonlinear_poisson exp_name = 'conv_mixed_residual_nonlinear' from models.darcy import conv_constitutive_constraint_nonlinear as constitutive_constraint hyparams = hyparams + f'_alpha1_{args.alpha1}_alpha2_{args.alpha2}' else: exp_name = 'conv_mixed_residual' from models.darcy import conv_constitutive_constraint as constitutive_constraint run_dir = args.exp_dir + '/' + exp_name + '/' + hyparams mkdirs(run_dir) # load data assert args.idx < 1000 if args.data == 'grf': assert args.kle in [512, 128, 1024, 2048] ntest = 1000 if args.kle == 512 else 1024 hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/kle{args.kle}_lhs{ntest}_test.hdf5' elif args.data == 'warped_grf': hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/warped_gp_ng64_n1000.hdf5' elif args.data == 'channelized': hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/channel_ng64_n512_test.hdf5' else: raise ValueError('No dataset are found for the speficied parameters') print(f'dataset: {hdf5_file}') with h5py.File(hdf5_file, 'r') as f: input_data = f['input'][()] output_data = f['output'][()] print(f'input: {input_data.shape}') print(f'output: {output_data.shape}') # permeability, (1, 1, 64, 64) perm_arr = input_data[[args.idx]] # pressure, flux_hor, flux_ver, (3, 64, 64) if args.nonlinear: # solve nonlinear Darcy for perm_arr with FEniCS output_file = run_dir + '/output_fenics.npy' if os.path.isfile(output_file): output_arr = np.load(output_file) print('Loaded solved output field') else: print('Solve nonlinear poisson with FEniCS...') output_arr = solve_nonlinear_poisson(perm_arr[0, 0], args.alpha1, args.alpha2, run_dir) np.save(output_file, output_arr) else: output_arr = output_data[args.idx] print('output shape: ', output_arr.shape) # model model = Decoder(args.nz, out_channels=3, blocks=args.blocks).to(device) print(f'model size: {model.model_size}') fixed_latent = torch.randn(1, args.nz, 16, 16).to(device) * 0.5 perm_tensor = torch.FloatTensor(perm_arr).to(device) sobel_filter = SobelFilter(args.imsize, correct=True, device=device) optimizer = optim.LBFGS(model.parameters(), lr=args.lr, max_iter=20, history_size=50) logger = {} logger['loss'] = [] def train(epoch): model.train() def closure(): optimizer.zero_grad() output = model(fixed_latent) if args.nonlinear: energy = constitutive_constraint(perm_tensor, output, sobel_filter, args.alpha1, args.alpha2) \ + continuity_constraint(output, sobel_filter) else: energy = constitutive_constraint( perm_tensor, output, sobel_filter) + continuity_constraint( output, sobel_filter) loss_dirichlet, loss_neumann = boundary_condition(output) loss_boundary = loss_dirichlet + loss_neumann loss = energy + loss_boundary * args.weight_bound loss.backward() if args.verbose: print(f'epoch {epoch}: loss {loss.item():6f}, '\ f'energy {energy.item():.6f}, diri {loss_dirichlet.item():.6f}, '\ f'neum {loss_neumann.item():.6f}') return loss loss = optimizer.step(closure) loss_value = loss.item() if not isinstance(loss, float) else loss logger['loss'].append(loss_value) print(f'epoch {epoch}: loss {loss_value:.6f}') if epoch % args.ckpt_freq == 0: torch.save(model.state_dict(), run_dir + "/model_epoch{}.pth".format(epoch)) def test(epoch): if epoch % args.epochs == 0 or epoch % args.test_freq == 0: output = model(fixed_latent) output = to_numpy(output) if args.animate: i_plot = epoch // args.test_freq plot_prediction_det_animate2(run_dir, output_arr, output[0], epoch, args.idx, i_plot, plot_fn='imshow', cmap=args.cmap, same_scale=args.same_scale) else: plot_prediction_det(run_dir, output_arr, output[0], epoch, args.idx, plot_fn='imshow', cmap=args.cmap, same_scale=args.same_scale) np.save(run_dir + f'/epoch{epoch}.npy', output[0]) print('start training...') dryrun = False tic = time.time() for epoch in range(1, args.epochs + 1): if not dryrun: train(epoch) test(epoch) print( f'Finished optimization for {args.epochs} epochs using {(time.time()-tic)/60:.3f} minutes' ) save_stats(run_dir, logger, 'loss') # save input plt.imshow(perm_arr[0, 0]) plt.colorbar() plt.savefig(run_dir + '/input.png') plt.close()
def main(): parser = argparse.ArgumentParser(description='CNN to solve PDE') parser.add_argument('--exp-dir', type=str, default='./experiments/solver', help='color map') # data parser.add_argument('--data-dir', type=str, default="./datasets", help='directory to dataset') parser.add_argument('--data', type=str, default='grf', choices=['grf', 'channelized', 'warped_grf'], help='data type') parser.add_argument('--kle', type=int, default=512, help='# kle terms') parser.add_argument('--imsize', type=int, default=64, help='image size') parser.add_argument('--idx', type=int, default=8, help='idx of input, please use 0 ~ 999') parser.add_argument('--alpha1', type=float, default=1.0, help='coefficient for the squared term') parser.add_argument('--alpha2', type=float, default=1.0, help='coefficient for the cubic term') parser.add_argument('--dim-hidden', type=int, default=512, help='# nodes in each hidden layer') parser.add_argument('--layers-hidden', type=int, default=8, help='# hidden layers') parser.add_argument('--off-grid', action='store_true', help='set True to use colloc ') parser.add_argument('--n-colloc', type=int, default=4096, help='# collocation points') parser.add_argument('--weight-bound', type=float, default=10, help='weight for boundary condition loss') parser.add_argument('--lr', type=float, default=0.5, help='learning rate') parser.add_argument('--epochs', type=int, default=2000, help='# epochs to train') parser.add_argument('--test-freq', type=int, default=50, help='every # epoch to test') parser.add_argument('--ckpt-freq', type=int, default=250, help='every # epoch to save model') parser.add_argument('--cmap', type=str, default='jet', help='color map') parser.add_argument('--same-scale', action='store_true', help='true for setting noise to be same scale as output') parser.add_argument('--animate', action='store_true', help='true to plot animate figures') parser.add_argument('--cuda', type=int, default=2, help='cuda number') parser.add_argument('-v', '--verbose', action='store_true', help='True for versbose output') args = parser.parse_args() pprint(vars(args)) device = torch.device(f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu") exp_name = 'fc_mixed_residual' dataset = f'{args.data}_kle{args.kle}' if args.data == 'grf' else args.data hyparams = f'{dataset}_idx{args.idx}_dhid{args.dim_hidden}_lhid{args.layers_hidden}_alpha1_{args.alpha1}_alpha2_{args.alpha2}_'\ f'lr{args.lr}_wb{args.weight_bound}_epochs{args.epochs}_ongrid_{not args.off_grid}_ncolloc{args.n_colloc}' run_dir = args.exp_dir + '/' + exp_name + '/' + hyparams mkdirs(run_dir) # load data assert args.idx < 1000 if args.data == 'grf': assert args.kle in [512, 128, 1024, 2048] ntest = 1000 if args.kle == 512 else 1024 hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/kle{args.kle}_lhs{ntest}_test.hdf5' elif args.data == 'warped_grf': hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/warped_gp_ng64_n1000.hdf5' elif args.data == 'channelized': assert args.idx < 512 hdf5_file = args.data_dir + f'/{args.imsize}x{args.imsize}/channel_ng64_n512_test.hdf5' else: raise ValueError('No dataset are found for the speficied parameters') print(f'dataset: {hdf5_file}') with h5py.File(hdf5_file, 'r') as f: input_data = f['input'][()] output_data = f['output'][()] print(f'input: {input_data.shape}') print(f'output: {output_data.shape}') # permeability, (1, 1, 64, 64) perm_arr = input_data[[args.idx]] # pressure, flux_hor, flux_ver, (3, 64, 64) output_arr = output_data[args.idx] def to_tensor_gpu(*numpy_seq): # x: numpy array --> tensor on GPU return (torch.FloatTensor(x).to(device) for x in numpy_seq) # define networks net_u = CPPN(dim_in=2, dim_out=3, dim_hidden=args.dim_hidden, layers_hidden=args.layers_hidden).to(device) print(net_u) print(net_u._model_size()) optimizer = optim.LBFGS(net_u.parameters(), lr=args.lr, max_iter=20, history_size=50) logger = {} logger['loss'] = [] ngrids = [args.imsize, args.imsize] sampler = SampleSpatial2d(int(ngrids[0]), int(ngrids[1])) colloc_on_grid = not args.off_grid # for batch optimization x_colloc = sampler.colloc(colloc_on_grid, n_samples=args.n_colloc).to(device) x_dirichlet = torch.cat((sampler.left(on_grid=False, n_samples=256), sampler.right(on_grid=False, n_samples=256)), 0).to(device) y_dirichlet = torch.cat((torch.ones(256, 1), torch.zeros(256, 1)), 0).to(device) x_neumann = torch.cat((sampler.top(colloc_on_grid), sampler.bottom(colloc_on_grid)), 0).to(device) print(sampler.coordinates_no_boundary.shape) K_true_tensor, = to_tensor_gpu(perm_arr.reshape(-1, 1)) if args.verbose: print('x_colloc: {}'.format(x_colloc.shape)) print('x_dirc: {}'.format(x_dirichlet.shape)) print('y_dirc: {}'.format(y_dirichlet.shape)) print('x_neumann: {}'.format(x_neumann.shape)) def train(epoch): net_u.train() def closure(): optimizer.zero_grad() loss_colloc = mixed_residual_fc(net_u, x_colloc, K_true_tensor, args.verbose, rand_colloc=args.off_grid) # loss_colloc = 0 loss_dirichlet = F.mse_loss(net_u(x_dirichlet)[:, [0]], y_dirichlet) loss_neumann = neumann_boundary_mixed(net_u, x_neumann) loss = loss_colloc + args.weight_bound * (loss_dirichlet + loss_neumann) loss.backward() if args.verbose: print(f'epoch {epoch}: colloc {loss_colloc:.6f}, ' f'diri {loss_dirichlet:.6f}, neum {loss_neumann:.6f}') return loss loss = optimizer.step(closure) loss_value = loss.item() if not isinstance(loss, float) else loss logger['loss'].append(loss_value) print('epoch {}: loss {:.10f}'.format(epoch, loss_value)) if epoch % args.ckpt_freq == 0: torch.save(net_u.state_dict(), run_dir + "/model_epoch{}.pth".format(epoch)) def test(epoch): if epoch % args.epochs == 0 or epoch % args.test_freq == 0: # plot the solution xx, yy = np.meshgrid(np.arange(ngrids[0]), np.arange(ngrids[1])) x_test = xx.flatten()[:, None] / ngrids[1] y_test = yy.flatten()[:, None] / ngrids[0] x_test, y_test = to_tensor_gpu(x_test, y_test) net_u.eval() x_test.requires_grad = True y_test.requires_grad = True xy_test = torch.cat((y_test, x_test), 1) y_pred = net_u(xy_test) target = output_arr # three output of net_u from 0-3 channel: u, flux_y, flux_x u_pred = y_pred[:, 0].detach().cpu().numpy().reshape(*ngrids) u_y = y_pred[:, 1].detach().cpu().numpy().reshape(*ngrids) u_x = y_pred[:, 2].detach().cpu().numpy().reshape(*ngrids) prediction = np.stack((u_pred, u_x, u_y)) # prediction = y_pred.view(*ngrids, -1).transpose(0, 1).permute(2, 1, 0).detach().cpu().numpy() if args.animate: i_plot = epoch // args.test_freq plot_prediction_det_animate2(run_dir, target, prediction, epoch, args.idx, i_plot, plot_fn='imshow', cmap=args.cmap, same_scale=args.same_scale) else: plot_prediction_det(run_dir, target, prediction, epoch, args.idx, plot_fn='imshow', cmap=args.cmap, same_scale=args.same_scale) np.save(run_dir + f'/epoch{epoch}.npy', prediction) print('start training...') dryrun = False tic = time.time() for epoch in range(1, args.epochs + 1): if not dryrun: train(epoch) test(epoch) print(f'Finished training {args.epochs} epochs in {(time.time()-tic)/60:.3f} minutes') save_stats(run_dir, logger, 'loss') # save input plt.close() plt.imshow(np.log(perm_arr[0, 0])) plt.colorbar() plt.savefig(run_dir + '/input_logK.png') plt.close() # super-resultion one ngrids = (640, 640) xx, yy = np.meshgrid(np.arange(ngrids[0]), np.arange(ngrids[1])) x_test = torch.FloatTensor(np.stack((yy.flatten() / (ngrids[1]-1), xx.flatten() / (ngrids[0]-1)), 1)).to(device) net_u.eval() u_pred = net_u(x_test) u_pred = u_pred[:, 0].reshape(*ngrids).detach().cpu().numpy() plt.contourf(u_pred, 65) plt.colorbar() plt.savefig(run_dir + '/solution_HR.png') plt.close()