def parse(self):
args = self.parse_args()
date = '{}'.format(time.strftime('%b_%d'))
args.run_dir = args.exp_dir + '/' + args.exp_name + '/' + date \
+ f'/{args.net}_noise_train{args.noise_levels_train}_'\
f'test{args.noise_levels_test}_{args.transform}_'\
f'epochs{args.epochs}_bs{args.batch_size}_lr{args.lr}'
args.ckpt_dir = args.run_dir + '/checkpoints'
if not args.post:
mkdirs([args.run_dir, args.ckpt_dir])
# seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.backends.cudnn.benchmark = True
print('Arguments:')
pprint(vars(args))
if not args.post:
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
def parse(self):
args = self.parse_args()
args.run_dir = args.exp_dir + '/' + args.exp_name \
+ '/kle{}/ntrain{}_blocks{}_growth{}_nif{}_drop{}_batch{}_lr{}_wd{}_epochs{}'.format(
args.kle, args.ntrain, args.blocks, args.growth_rate,
args.init_features, args.drop_rate, args.batch_size,
args.lr, args.weight_decay, args.epochs
)
args.ckpt_dir = args.run_dir + '/checkpoints'
mkdirs([args.run_dir, args.ckpt_dir])
assert args.epochs % args.ckpt_freq == 0, 'epochs must'\
'be dividable by ckpt_freq'
# seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
print('Arguments:')
pprint(vars(args))
if not args.post:
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
def parse(self):
args = self.parse_args()
args.run_dir = args.exp_dir + '/' + args.exp_name \
+ '/kle{}/nsamples{}_ntrain{}_batch{}_lr{}_noiselr{}_epochs{}'.format(
args.kle, args.n_samples, args.ntrain, args.batch_size, args.lr,
args.lr_noise, args.epochs)
args.ckpt_dir = args.run_dir + '/checkpoints'
mkdirs([args.run_dir, args.ckpt_dir])
assert args.epochs % args.ckpt_freq == 0, 'epochs must'\
'be dividable by ckpt_freq'
assert args.ntrain % args.batch_size == 0, 'num of training data must'\
'be dividable by batch size'
# seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
print('Arguments:')
pprint(vars(args))
if not args.post:
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
def parse(self, save=True):
if not self.initialized:
self.initialize()
self.opt = self.parser.parse_args()
self.opt.isTrain = self.isTrain # train or test
str_ids = self.opt.gpu_ids.split(',')
self.opt.gpu_ids = []
for str_id in str_ids:
id = int(str_id)
if id >= 0:
self.opt.gpu_ids.append(id)
# set gpu ids
if len(self.opt.gpu_ids) > 0:
torch.cuda.set_device(self.opt.gpu_ids[0])
args = vars(self.opt)
print('------------ Options -------------')
for k, v in sorted(args.items()):
print('%s: %s' % (str(k), str(v)))
print('-------------- End ----------------')
# save to the disk
expr_dir = os.path.join(self.opt.checkpoints_dir, self.opt.name)
misc.mkdirs(expr_dir)
if save and not self.opt.continue_train:
file_name = os.path.join(expr_dir, 'opt.txt')
with open(file_name, 'wt') as opt_file:
opt_file.write('------------ Options -------------\n')
for k, v in sorted(args.items()):
opt_file.write('%s: %s\n' % (str(k), str(v)))
opt_file.write('-------------- End ----------------\n')
return self.opt
def parse(self):
args = self.parse_args()
tid = 2
hparams = f'{args.data}_{tid}_run{args.run}_bs{args.batch_size}'
# if args.debug:
# hparams = 'debug/' + hparams
args.run_dir = args.exp_dir + '/' + args.exp_name + '/' + hparams
args.ckpt_dir = args.run_dir + '/checkpoints'
# print(args.run_dir)
# print(args.ckpt_dir)
mkdirs(args.run_dir, args.ckpt_dir)
# assert args.ntrain % args.batch_size == 0 and \
# args.ntest % args.test_batch_size == 0
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
print('Arguments:')
pprint(vars(args))
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
def _save(self, args):
expr_dir = os.path.join(self._opt.checkpoints_dir, self._opt.name)
print(expr_dir)
misc.mkdirs(expr_dir)
file_name = os.path.join(expr_dir, 'opt_%s.txt' % ('train' if self.is_train else 'test'))
with open(file_name, 'wt') as opt_file:
opt_file.write('------------ Options -------------\n')
for k, v in sorted(args.items()):
opt_file.write('%s: %s\n' % (str(k), str(v)))
opt_file.write('-------------- End ----------------\n')
def parse(self):
args = self.parse_args()
args.LU_decompose = not args.no_LU_decompose
assert len(args.enc_blocks) == len(args.flow_blocks)
hparams = f'kle{args.kle}_ntrain{args.ntrain}_'\
f'ENC_blocks{args.enc_blocks}_FLOW_blocks{args.flow_blocks}_'\
f'wb{args.weight_bound}_beta{args.beta}_'\
f'batch{args.batch_size}_lr{args.lr}_epochs{args.epochs}'
if args.debug:
hparams = 'debug/' + hparams
if args.data_init:
hparams = hparams + '_data_init'
args.run_dir = args.exp_dir + '/' + args.exp_name + '/' + hparams
args.ckpt_dir = args.run_dir + '/checkpoints'
args.train_dir = args.run_dir + '/training'
args.pred_dir = args.train_dir + '/predictions'
mkdirs(args.run_dir, args.ckpt_dir, args.train_dir, args.pred_dir)
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random Seed: ", args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
args_file = args.run_dir + "/args.txt"
if os.path.isfile(args_file):
# args.resume, args.ckpt directly overwrite
if args.ckpt_epoch is None and args.resume:
with open(args_file, 'r') as args_f:
args_old = argparse.Namespace(**json.load(args_f))
args.ckpt_epoch = args_old.ckpt_epoch
else:
with open(args_file, 'w') as args_f:
json.dump(vars(args), args_f, indent=4)
print('Arguments:')
pprint(vars(args))
return args
pprint(vars(args))
if not args.post:
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
args = Parser().parse()
device = torch.device(
f'cuda:{args.cuda}' if torch.cuda.is_available() else 'cpu')
args.train_dir = args.run_dir + "/training"
args.pred_dir = args.train_dir + "/predictions"
mkdirs([args.train_dir, args.pred_dir])
if args.net == 'unet':
model = UnetN2N(args.in_channels, args.out_channels).to(device)
elif args.net == 'unetv2':
model = UnetN2Nv2(args.in_channels, args.out_channels).to(device)
if args.debug:
print(model)
print(model.model_size)
if args.transform == 'four_crop':
# wide field images may have complete noise in center-crop case
transform = transforms.Compose([
transforms.FiveCrop(args.imsize),
transforms.Lambda(lambda crops: torch.stack(
[fluore_to_tensor(crop) for crop in crops[:4]])),
print('Arguments:')
pprint(vars(args))
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
args = Parser().parse()
device = torch.device(
f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
print(device)
args.train_dir = args.run_dir + '/training'
args.pred_dir = args.train_dir + '/predictions'
mkdirs(args.train_dir, args.pred_dir)
data, ref = LoadData(device)
model = DenseED(in_channels=1,
out_channels=2,
imsize=args.imsize,
blocks=args.blocks,
growth_rate=args.growth_rate,
init_features=args.init_features,
drop_rate=args.drop_rate,
out_activation=None,
upsample=args.upsample).to(device)
# modelname = "experiments/codec/mixed_residual/debug/SIMP_penal3_plus1_1_run1_bs64/checkpoints/model_epoch5300.pth"
# model = torch.load(modelname)
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()
with open(args.run_dir + "/args.txt", 'w') as args_file:
json.dump(vars(args), args_file, indent=4)
return args
if __name__ == '__main__':
args = Parser().parse()
device = torch.device(
f"cuda:{args.cuda}" if torch.cuda.is_available() else "cpu")
args.train_dir = args.run_dir + '/training'
args.pred_dir = args.train_dir + '/predictions'
mkdirs(args.train_dir, args.pred_dir)
model = DenseED(in_channels=1,
out_channels=3,
imsize=args.imsize,
blocks=args.blocks,
growth_rate=args.growth_rate,
init_features=args.init_features,
drop_rate=args.drop_rate,
out_activation=None,
upsample=args.upsample).to(device)
if args.debug:
print(model)
# if start from ckpt
if args.ckpt_epoch is not None:
ckpt_file = args.run_dir + f'/checkpoints/model_epoch{args.ckpt_epoch}.pth'
parser.add_argument('--cuda', type=int, default=0, help='gpu #')
args_post = parser.parse_args()
run_dir = './experiments/cglow/reverse_kld/kle100_ntrain4096_ENC_blocks[3, 4, 4]_'\
'FLOW_blocks[6, 6, 6]_wb50_beta150.0_batch32_lr0.0015_epochs400'
if args_post.run_dir is not None:
run_dir = args_post.run_dir
device = torch.device(f'cuda:{args_post.cuda}' if torch.cuda.is_available() else 'cpu')
# load the args for pre-trained run
args = load_args(run_dir)
args.device = device
args.post_dir = args.run_dir + f'/post_proc_mc{args_post.n_mc}_nsamples{args_post.n_samples}_'\
f'varsamples{args_post.var_samples}_temp{args_post.temperature}'
mkdirs(args.post_dir)
# load the pre-trained model
cglow = MultiScaleCondGlow(img_size=args.imsize,
x_channels=args.x_channels,
y_channels=args.y_channels,
enc_blocks=args.enc_blocks,
flow_blocks=args.flow_blocks,
LUdecompose=args.LU_decompose,
squeeze_factor=2,
data_init=args.data_init)
print(cglow.model_size)
ckpt_file = args.ckpt_dir + f"/model_epoch{args.epochs}.pth"
checkpoint = torch.load(ckpt_file, map_location='cpu')
cglow.load_state_dict(checkpoint['model_state_dict'])
cglow = cglow.to(device)
]
image_types = [
'TwoPhoton_BPAE_R', 'TwoPhoton_BPAE_G', 'TwoPhoton_BPAE_B',
'TwoPhoton_MICE'
]
# image_types = ['test_mix']
run_dir = args_test.pretrain_dir + f'/{args_test.model}'
with open(run_dir + '/args.txt') as args_file:
args = Namespace(**json.load(args_file))
pprint(args)
if args_test.no_cuda:
test_dir = run_dir + '/benchmark_cpu'
else:
test_dir = run_dir + '/benchmark_gpu'
mkdirs(test_dir)
if args_test.model == 'dncnn':
model = DnCNN(depth=args.depth,
n_channels=args.width,
image_channels=1,
use_bnorm=True,
kernel_size=3)
elif args_test.model == 'n2n':
model = UnetN2N(args.in_channels, args.out_channels)
if args.debug:
print(model)
print(module_size(model))
model.load_state_dict(
torch.load(run_dir + f'/checkpoints/model_epoch{args.epochs}.pth',
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()
noise_levels = [1]
image_types = [
'Confocal_BPAE_R', 'Confocal_BPAE_G', 'Confocal_BPAE_B', 'Confocal_FISH'
]
data_dir = args_test.data_root
run_dir = args_test.pretrain_dir + f'/{args_test.model}'
with open(run_dir + '/args.txt') as args_file:
args = Namespace(**json.load(args_file))
pprint(args)
if args_test.no_cuda:
test_dir = run_dir + '/example_cpu'
else:
test_dir = run_dir + '/example_gpu'
mkdirs(test_dir)
if args_test.model == 'dncnn':
model = DnCNN(depth=args.depth,
n_channels=args.width,
image_channels=1,
use_bnorm=True,
kernel_size=3)
elif args_test.model == 'n2n':
model = UnetN2N(args.in_channels, args.out_channels)
if args.debug:
print(model)
print(module_size(model))
model.load_state_dict(
torch.load(run_dir + f'/checkpoints/model_epoch{args.epochs}.pth',
