From 78c48745ac232f204b4416088a1bdd1b37800291 Mon Sep 17 00:00:00 2001 From: Sourcery AI <> Date: Fri, 11 Nov 2022 19:29:38 +0000 Subject: [PATCH] 'Refactored by Sourcery' --- learner.py | 44 +++++++++++++++++++++---------------------- simulation_dataset.py | 26 ++++++++++++------------- waveynet_trainer.py | 27 +++++++++++++++++--------- 3 files changed, 53 insertions(+), 44 deletions(-) diff --git a/learner.py b/learner.py index 4f961c4..f4a9374 100644 --- a/learner.py +++ b/learner.py @@ -15,10 +15,10 @@ def __init__(self, args): ''' self.alpha = args.alpha config = [] + # kernel_size = 2 if block==args.num_down_conv-1 else 3 + kernel_size = 3 for block in range(args.num_down_conv): - # kernel_size = 2 if block==args.num_down_conv-1 else 3 - kernel_size = 3 out_channels = (2**block)*args.hidden_dim if(block == 0): config += [ @@ -104,7 +104,7 @@ def __init__(self, args): self.lr_scheduler = None # will be initialized in waveynet_trainer.py self.residual_terms = None # to store the residual connect for addition later - for i, (name, param) in enumerate(self.config): + for name, param in self.config: if name == 'conv2d': # [ch_out, ch_in, kernelsz, kernelsz] w = nn.Parameter(torch.ones(*param[:4]), requires_grad=True) @@ -144,12 +144,12 @@ def extra_repr(self): for name, param in self.config: if name == 'conv2d': tmp = 'conv2d:(ch_in:%d, ch_out:%d, k:%dx%d, stride:%d, padding:%d)'\ - %(param[1], param[0], param[2], param[3], param[4], param[5],) + %(param[1], param[0], param[2], param[3], param[4], param[5],) info += tmp + '\n' elif name == 'conv2d_b': tmp = 'conv2d:(ch_in:%d, ch_out:%d, k:%dx%d, stride:%d, padding:%d)'\ - %(param[1], param[0], param[2], param[3], param[4], param[5],) + %(param[1], param[0], param[2], param[3], param[4], param[5],) info += tmp + '\n' elif name == 'leakyrelu': @@ -160,8 +160,8 @@ def extra_repr(self): tmp = 'max_pool2d:(k:%s, stride:%s, padding:%d)'%(param[0], param[1], param[2]) info += tmp + '\n' elif name in ['flatten', 'tanh', 'relu', 'upsample', 'reshape', 'sigmoid',\ - 'use_logits', 'bn', 'residual']: - tmp = name + ':' + str(tuple(param)) + 'use_logits', 'bn', 'residual']: + tmp = f'{name}:{tuple(param)}' info += tmp + '\n' else: raise NotImplementedError @@ -183,7 +183,14 @@ def forward(self, x, vars=None, bn_training=True): blocks = [] for name, param in self.config: - if name == 'conv2d': + if name == 'bn': + w, b = vars[idx], vars[idx + 1] + running_mean, running_var = self.vars_bn[bn_idx], self.vars_bn[bn_idx+1] + x = F.batch_norm(x, running_mean, running_var, weight=w, bias=b, training=bn_training) + idx += 2 + bn_idx += 2 + + elif name == 'conv2d': # periodic padding kernel_width = param[2] if kernel_width % 2: @@ -222,16 +229,16 @@ def forward(self, x, vars=None, bn_training=True): x = F.conv2d(x_pad, w, b, stride=param[4], padding=0) idx += 2 - elif name == 'bn': - w, b = vars[idx], vars[idx + 1] - running_mean, running_var = self.vars_bn[bn_idx], self.vars_bn[bn_idx+1] - x = F.batch_norm(x, running_mean, running_var, weight=w, bias=b, training=bn_training) - idx += 2 - bn_idx += 2 - elif name == 'leakyrelu': x = F.leaky_relu(x, negative_slope=param[0], inplace=param[1]) + elif name == 'max_pool2d': + blocks.append(x) + x = F.max_pool2d(x, param[0], stride=param[1], padding=param[2]) + + elif name == 'residual': + x = x + self.residual_terms + elif name == 'upsample': if first_upsample: first_upsample = False @@ -240,13 +247,6 @@ def forward(self, x, vars=None, bn_training=True): x = F.interpolate(x, size=(shortcut.shape[2],shortcut.shape[3]), mode='nearest') x = torch.cat([shortcut,x], dim=1) # batch, channels, h, w - elif name == 'residual': - x = x + self.residual_terms - - elif name == 'max_pool2d': - blocks.append(x) - x = F.max_pool2d(x, param[0], stride=param[1], padding=param[2]) - else: print(name) raise NotImplementedError diff --git a/simulation_dataset.py b/simulation_dataset.py index 7078fbe..f03b9fe 100644 --- a/simulation_dataset.py +++ b/simulation_dataset.py @@ -21,18 +21,18 @@ def __init__(self, data_folder, local_data, total_sample_number = None, transfor print("Loading training data...") print("Note: The training and test datasets are very large, which can take "\ - "on the order of hours to complete, depending on your internet download "\ - "speed. If you have a slower internet download speed, consider downloading the "\ - "data locally from http://metanet.stanford.edu/search/waveynet-study/ so "\ - "that you do not have to re-download the data each time you run the code."\ - "See the README file for more information.") + "on the order of hours to complete, depending on your internet download "\ + "speed. If you have a slower internet download speed, consider downloading the "\ + "data locally from http://metanet.stanford.edu/search/waveynet-study/ so "\ + "that you do not have to re-download the data each time you run the code."\ + "See the README file for more information.") #load the training data, either locally or from Metanet, as specified by #the boolean args.local_data - if(local_data): - train_data = np.load(data_folder+'/train_ds.npz') + if local_data: + train_data = np.load(f'{data_folder}/train_ds.npz') else: response = requests.get('http://metanet.stanford.edu/static/search/waveynet/'\ - 'data/train_ds.npz') + 'data/train_ds.npz') response.raise_for_status() train_data = np.load(io.BytesIO(response.content)) @@ -40,11 +40,11 @@ def __init__(self, data_folder, local_data, total_sample_number = None, transfor #load the test data, either locally or from Metanet, as specified by #the boolean args.local_data - if(local_data): - test_data = np.load(data_folder+'/test_ds.npz') + if local_data: + test_data = np.load(f'{data_folder}/test_ds.npz') else: response = requests.get('http://metanet.stanford.edu/static/search/waveynet/'\ - 'data/test_ds.npz') + 'data/test_ds.npz') response.raise_for_status() test_data = np.load(io.BytesIO(response.content)) @@ -72,9 +72,9 @@ def __init__(self, data_folder, local_data, total_sample_number = None, transfor for dev in range(len(self.input_imgs)): self.input_imgs[dev][np.where(self.input_imgs[dev]==self.input_imgs[dev,0,0,0])]=\ - die_change[dev] + die_change[dev] self.input_imgs[dev]=(self.input_imgs[dev]*(np.sqrt(self.refr_idx_vec[dev]) - \ - n_air) + n_air)**2 + n_air) + n_air)**2 #Use only a subset of the dataset if specified to do so by args.total_sample_number if total_sample_number: diff --git a/waveynet_trainer.py b/waveynet_trainer.py index ecd73af..6cd69a7 100644 --- a/waveynet_trainer.py +++ b/waveynet_trainer.py @@ -147,8 +147,15 @@ def H_to_H(Hz_R: np.array, Hz_I: np.array, dL: float, omega: float, eps_grid: np FD_Ex = Hz_to_Ex(Hz_R, Hz_I, dL, omega, eps_grid, eps_0) FD_Ez = Hz_to_Ey(Hz_R, Hz_I, dL, omega, eps_grid, eps_0) - FD_H = E_to_Hz(FD_Ez[:, 0, :-1], FD_Ez[:, 1, :-1], FD_Ex[:, 0], FD_Ex[:, 1], dL, omega, mu_0) - return FD_H + return E_to_Hz( + FD_Ez[:, 0, :-1], + FD_Ez[:, 1, :-1], + FD_Ex[:, 0], + FD_Ex[:, 1], + dL, + omega, + mu_0, + ) def regConstScheduler(epoch, args, last_epoch_data_loss, last_epoch_physical_loss): ''' @@ -216,14 +223,16 @@ def main(args): test_loader = DataLoader(test_ds, batch_size=args.batch_size, shuffle=True, num_workers=0) - train_mean = 0 - test_mean = 0 + train_mean = sum( + torch.mean(torch.abs(sample_batched["field"])) + for sample_batched in train_loader + ) + + test_mean = sum( + torch.mean(torch.abs(sample_batched["field"])) + for sample_batched in test_loader + ) - # first get the mean-absolute-field value: - for sample_batched in train_loader: - train_mean += torch.mean(torch.abs(sample_batched["field"])) - for sample_batched in test_loader: - test_mean += torch.mean(torch.abs(sample_batched["field"])) train_mean /= len(train_loader) test_mean /= len(test_loader)