def _cam_image(self, norm_image): """Get CAM for an image. Args: norm_image: Normalized image. Should be of type torch.Tensor or a dictionary containing image, target and prediction Returns: Dictionary containing unnormalized image, heatmap, CAM result, target and pred """ if type(norm_image) == dict: target = norm_image['target'] pred = norm_image['pred'] norm_image = norm_image['img'].cpu() else: target = pred = "" norm_image_cuda = norm_image.clone().unsqueeze_(0).to(self.device) heatmap, result = {}, {} for layer, gc in self.gradcam.items(): mask, _ = gc(norm_image_cuda) cam_heatmap, cam_result = visualize_cam( mask, unnormalize(norm_image, self.mean, self.std, out_type='tensor').clone().unsqueeze_(0).to( self.device), alpha=self.heatmap_alpha) heatmap[layer], result[layer] = to_numpy(cam_heatmap), to_numpy( cam_result) return { 'image': unnormalize(norm_image, self.mean, self.std), 'heatmap': heatmap, 'result': result, 'target': target, 'pred': pred }
def _train_step(next_lr, input2, input3, input4, hr2, hr3, hr4): # print ('_train_step') opt.learning_rate.assign(next_lr) with tf.GradientTape() as tape: output2 = net(input2) output3 = net(input3) output4 = net(input4) loss2 = tf.reduce_mean(tf.abs(output2 - normalize(hr2))) loss3 = tf.reduce_mean(tf.abs(output3 - normalize(hr3))) loss4 = tf.reduce_mean(tf.abs(output4 - normalize(hr4))) cost = loss2 + loss3 + loss4 # tf.print(loss2) grads = tape.gradient(cost, net.trainable_variables) opt.apply_gradients(zip(grads, net.trainable_variables)) y_pred2 = clip255(unnormalize(output2)) y_pred3 = clip255(unnormalize(output3)) y_pred4 = clip255(unnormalize(output4)) db2 = tf.reduce_mean(tf.math.minimum(tf.image.psnr(y_pred2, hr2, 255), 100)) db3 = tf.reduce_mean(tf.math.minimum(tf.image.psnr(y_pred3, hr3, 255), 100)) db4 = tf.reduce_mean(tf.math.minimum(tf.image.psnr(y_pred4, hr4, 255), 100)) return cost, db2, db3, db4
def _dynamics_func(self, state, action, reuse=True): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ print("ModelBasedPolicy - before _dynamics_func") ### PROBLEM 1 ############## ### part a ###Normalize state and action by using statistics of self._init_dataset ############## s_mean = self._init_dataset.state_mean a_mean = self._init_dataset.action_mean s_std = self._init_dataset.state_std a_std = self._init_dataset.action_std state_d_m = self._init_dataset.delta_state_mean state_d_s = self._init_dataset.delta_state_std normalize_state = utils.normalize(state, s_mean, s_std) normalize_action = utils.normalize(action, a_mean, a_std) ############## ### part b ### Concatenate state and action ############## s_a = tf.concat([normalize_state, normalize_action], axis=1) # s_a = np.concatenate((normalize_state,normalize_action),axis = None) # d = s_a.shape ############## ### part c ### Generate NN ############## # print ("generate the NN") # s_a_placeholder = tf.placeholder(tf.float32, [None,d[0]]) print("generate the NN") next_state_prediction = utils.build_mlp(s_a, self._state_dim, "nn_prediciton", n_layers=self._nn_layers, reuse=reuse) print("finish gen NN") ############## ### pard d ### ############## delta_state = utils.unnormalize(next_state_prediction, state_d_m, state_d_s) next_state_pred = state + delta_state print("ModelBasedPolicy - after _dynamics_func") return next_state_pred
def dynamics_func(self, state, action, reuse): """ takes state and action, returns the next state returns: prediction of next state """ state_norm = utils.normalize(state, self.init_dataset.state_mean, self.init_dataset.state_std) action_norm = utils.normalize(action, self.init_dataset.action_mean, self.init_dataset.action_std) # network input is concatenated state, action s_a = tf.concat([state_norm, action_norm], axis=1) d_next_state_norm = utils.build_mlp(s_a, self.state_dim, 'prediction', self.nn_layers, reuse=reuse) d_next_state = utils.unnormalize(d_next_state_norm, self.init_dataset.delta_state_mean, self.init_dataset.delta_state_std) next_state_pred = d_next_state + state return next_state_pred
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE # REMEMBER: use a new variable such as 'normalized_state' instead of reuse variable 'state', OTHERWISE, it won't work (maybe because of no gradient update to it) normalized_state = utils.normalize(state, self._init_dataset.state_mean, self._init_dataset.state_std) normalized_action = utils.normalize(action, self._init_dataset.action_mean, self._init_dataset.action_std) normalized_state_dif = utils.build_mlp(tf.concat([normalized_state, normalized_action], axis=1), self._state_dim, scope="dynamics_func", n_layers=self._nn_layers, reuse=reuse) next_state_pred = state + utils.unnormalize(normalized_state_dif, self._init_dataset.delta_state_mean, self._init_dataset.delta_state_std) return next_state_pred
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE state_norm = utils.normalize(state, self._init_dataset.state_mean, self._init_dataset.state_std) action_norm = utils.normalize(action, self._init_dataset.action_mean, self._init_dataset.action_std) x = tf.concat([state_norm, action_norm], axis=1) x = utils.build_mlp(x, self._state_dim, 'dynamic_net', self._nn_layers, reuse=reuse) next_state_pred = state + utils.unnormalize( x, self._init_dataset.state_mean, self._init_dataset.state_std) return next_state_pred
def visualize_weights(hdf5_file_name, ext='.png'): #with h5py.File(file_name, 'r') as file: image_folder_path = osp.join( osp.dirname(osp.abspath(__file__)), 'weights_visualization_' + osp.splitext(osp.basename(hdf5_file_name))[0]) if not (osp.exists(image_folder_path)): os.mkdir(image_folder_path) mat_dict = {} is_first_tensor = True for (path, item) in read_hdf5(hdf5_file_name): if path.startswith('/model_weights'): tensor = np.array(item) # print(tensor.shape) if len(tensor.shape) == 4: if is_first_tensor: wmin = np.amin(tensor) wmax = np.amax(tensor) is_first_tensor = False print(wmin) print(wmax) else: wmin = min(wmin, np.amin(tensor)) wmax = max(wmax, np.amax(tensor)) for j in range(tensor.shape[3]): for i in range(tensor.shape[2]): mat = tensor[:, :, i, j] #print(mat) image_file_name = (path[1:].replace('/', '_')).replace( ':', '_') + '_{}_{}'.format(i, j) + ext image_path = osp.join(image_folder_path, image_file_name) #print(image_path) mat_dict[image_path] = mat print(wmin) print(wmax) for image_path in mat_dict.keys(): mat_dict[image_path] = unnormalize(mat_dict[image_path], wmin, wmax) cv2.imwrite(image_path, mat_dict[image_path]) # for key in file.keys(): # print(key) # if(key == "model_weights"): # print("yes") # weights = file[key] # print(weights) #print(weights) # print("break") # break #print(min(data)) #print(max(data)) #print(data[:15]) return 0
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE # self._dynamics_func is supposed to take in a batch. The input state is assumed to be [None, self._state_dim]. state_mean, state_std = self._init_dataset.state_mean, self._init_dataset.state_std action_mean, action_std = self._init_dataset.action_mean, self._init_dataset.action_std normalized_state = utils.normalize(state, mean=state_mean, std=state_std) normalized_action = utils.normalize(action, mean=action_mean, std=action_std) input_nn = tf.concat(values=[normalized_state, normalized_action], axis=1) #tf.concat([normalized_state, normalized_action], axis=-1) # np.concatenate(normalized_state, normalized_action) # print("unnormlaized state: ", state) # print("state: ", normalized_state.shape) # print("state: ", normalized_state) # print("action: ", normalized_action.shape) # print("type: ", type(normalized_action)) # print("input_nn: : ", input_nn) normalized_delta_state_pred = utils.build_mlp( input_layer=input_nn, output_dim=self._state_dim, scope="dynamics_model", n_layers=self._nn_layers, reuse=reuse) delta_state_pred = utils.unnormalize( normalized_delta_state_pred, self._init_dataset.delta_state_mean, self._init_dataset.delta_state_std) next_state_pred = tf.add(state, delta_state_pred) return next_state_pred
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE ## (a) Normalize both the state and action by using the statistics of self._init_dataset and ## the utils.normalize function state_mean = self._init_dataset.state_mean state_std = self._init_dataset.state_std state_normalize = utils.normalize(state, state_mean, state_std, eps=1e-8) # # @@@@@@ # print("state", tf.convert_to_tensor(state).get_shape()) # print("action", tf.convert_to_tensor(action).get_shape()) # state (?, 20) # action (?, 6) action_mean = self._init_dataset.action_mean action_std = self._init_dataset.action_std action_normalize = utils.normalize(action, action_mean, action_std, eps=1e-8) ## (b) Concatenate the normalized state and action concatenated = tf.concat([state_normalize, action_normalize], axis=1, name='concatenated') # (c) Pass the concatenated, normalized state-action tensor through a neural network with # self._nn_layers number of layers using the function utils.build_mlp. The resulting output # is the normalized predicted difference between the next state and the current state next_state = utils.build_mlp(input_layer=concatenated, output_dim=self._state_dim, scope='dynamics_func', n_layers=self._nn_layers, reuse=reuse) # (d) Unnormalize the delta state prediction, and add it to the current state in order to produce # the predicted next state delta_state_mean = self._init_dataset.delta_state_mean delta_state_std = self._init_dataset.delta_state_std next_state_pred = state + utils.unnormalize( next_state, delta_state_mean, delta_state_std) return next_state_pred
def write_result(fs, pred_dict, attr): pred = pred_dict['pred'].data.cpu().numpy() label = pred_dict['label'].data.cpu().numpy() for i in range(pred_dict['pred'].size()[0]): # fs.write('%.6f %.6f\n' % (label[i][0], pred[i][0])) weekID = attr['weekID'].data[i] timeID = attr['timeID'].data[i] driverID = attr['driverID'].data[i] dist = utils.unnormalize(attr['dist'].data[i], 'dist') fs.write('%d,%d,%d,%.6f,%.6f,%.6f\n' % (weekID, timeID, driverID, dist, label[i][0], pred[i][0]))
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE #print(type(state) norm_state = utils.normalize(state, self._init_dataset.state_mean, self._init_dataset.state_std) norm_action = utils.normalize(action, self._init_dataset.action_mean, self._init_dataset.action_std) norm_all = tf.concat((norm_state, norm_action), axis=1) # st_ph,ac_ph,nx_st_ph = self._setup_placeholders() dy_func = utils.build_mlp(norm_all, self._state_dim, scope="DYN_FUNC", n_layers=self._nn_layers, reuse=reuse) # dy_out = dy_func(norm_all) #self.sess.run(dy_func,feed_dict={norm_all:norm_all}) dy_out = utils.unnormalize(dy_func, self._init_dataset.delta_state_mean, self._init_dataset.delta_state_std) next_state_pred = dy_out + state #raise NotImplementedError return next_state_pred
def validate(scale): dbs = [] for i in range(len(eval_gt_imgs)): hr = eval_gt_imgs[i] hr = modular_crop(hr, scale) h, w, _ = hr.shape lr = resize_img(hr, h // scale, w // scale) upcubic = resize_img(lr, h, w) inputs = np.asarray([lr], np.float32), np.asarray([upcubic], np.float32), scale output = net(inputs) y_pred = clip255(unnormalize(output)) dbs.append(psnr(y_pred, hr)) return np.mean(dbs)
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ normalized_state = utils.normalize(state, mean=self._init_dataset.state_mean, std=self._init_dataset.state_std) normalized_action = utils.normalize( action, mean=self._init_dataset.action_mean, std=self._init_dataset.action_std) normalized_state_action = tf.concat( [normalized_state, normalized_action], axis=1) normalized_state_delta = utils.build_mlp( input_layer=normalized_state_action, output_dim=self._state_dim, scope='dynamics_func', n_layers=self._nn_layers, reuse=reuse) state_delta = utils.unnormalize( normalized_state_delta, mean=self._init_dataset.delta_state_mean, std=self._init_dataset.delta_state_std) next_state_pred = state + state_delta return next_state_pred
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE # convert to float32 state = tf.dtypes.cast(state, dtype=tf.float32) action = tf.dtypes.cast(action, dtype=tf.float32) mean_state = self._init_dataset.state_mean std_state = self._init_dataset.state_std state_norm = utils.normalize(state, mean_state, std_state) mean_action = self._init_dataset.action_mean std_action = self._init_dataset.action_std action_norm = utils.normalize(action, mean_action, std_action) state_action = tf.concat([state_norm, action_norm], axis=1) delta_state_prediction_norm = utils.build_mlp(state_action, self._state_dim, 'state_prediction', n_layers=self._nn_layers, reuse=reuse) delta_mean_state = self._init_dataset.delta_state_mean delta_std_state = self._init_dataset.delta_state_std delta_state_prediction = utils.unnormalize(delta_state_prediction_norm, delta_mean_state, delta_std_state) next_state_pred = state + delta_state_prediction return next_state_pred
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE #print(self._init_dataset.shape) # Normalize state and action #d_mean, d_std = np.mean(self._init_dataset), np.std(self._init_dataset) n_state = utils.normalize(state, self._init_dataset.state_mean, self._init_dataset.state_std) n_action = utils.normalize(action, self._init_dataset.action_mean, self._init_dataset.action_std) # Predict next state based on the current state and action n_sa = tf.concat((n_state, n_action), axis=1) n_nsp = utils.build_mlp(n_sa, output_dim=self._state_dim, scope="f_transition", n_layers=self._nn_layers, reuse=reuse) next_state_pred = state + utils.unnormalize( n_nsp, self._init_dataset.delta_state_mean, self._init_dataset.delta_state_std) return next_state_pred
def save_images(self, X_source, X_target, images_dir, nb_images=64): images_dir = os.path.join(self.output_dir, "generated-images") if not os.path.exists(images_dir): os.mkdir(images_dir) X_s2s = unnormalize(self.sess.run(self.G_s2s, feed_dict={self.ipt_source: X_source[:nb_images]})) X_t2t = unnormalize(self.sess.run(self.G_t2t, feed_dict={self.ipt_target: X_target[:nb_images]})) X_s2t = unnormalize(self.sess.run(self.G_s2t, feed_dict={self.ipt_source: X_source[:nb_images]})) X_t2s = unnormalize(self.sess.run(self.G_t2s, feed_dict={self.ipt_target: X_target[:nb_images]})) X_cycle_s2s = unnormalize(self.sess.run(self.G_cycle_s2s, feed_dict={self.ipt_source: X_source[:nb_images]})) X_cycle_t2t = unnormalize(self.sess.run(self.G_cycle_t2t, feed_dict={self.ipt_target: X_target[:nb_images]})) Y_source_predict = self.sess.run(self.D_source_predict, feed_dict={self.ipt_source: X_source[:nb_images]}) Y_target_predict = self.sess.run(self.DG_t2s_predict, feed_dict={self.ipt_target: X_target[:nb_images]}) for index in range(len(X_s2s)): plot_images(index, X_source, X_target, X_s2s, X_t2t, X_s2t, X_t2s, X_cycle_s2s, X_cycle_t2t, Y_source_predict, Y_target_predict) print("Saving 'generated-images/iter_{:05d}_image_{:02d}.png'".format(self.iter, index), end="\r") plt.savefig(os.path.join(images_dir, "iter_{:05d}_image_{:02d}.png".format(self.iter, index)))
def _dynamics_func(self, state, action, reuse=False): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE # raise NotImplementedError norm_state = utils.normalize(state, self._init_dataset.state_mean, self._init_dataset.state_std) norm_action = utils.normalize(action, self._init_dataset.action_mean, self._init_dataset.action_std) norm_state_action = tf.concat([norm_state, norm_action], 1) norm_delta_state_pred = utils.build_mlp(norm_state_action, self._state_dim, "dynamic", n_layers=self._nn_layers, reuse=reuse) delta_state_pred = utils.unnormalize( norm_delta_state_pred, self._init_dataset.delta_state_mean, self._init_dataset.delta_state_std) next_state_pred = state + delta_state_pred return next_state_pred
def _dynamics_func(self, state, action, reuse): """ Takes as input a state and action, and predicts the next state returns: next_state_pred: predicted next state implementation details (in order): (a) Normalize both the state and action by using the statistics of self._init_dataset and the utils.normalize function (b) Concatenate the normalized state and action (c) Pass the concatenated, normalized state-action tensor through a neural network with self._nn_layers number of layers using the function utils.build_mlp. The resulting output is the normalized predicted difference between the next state and the current state (d) Unnormalize the delta state prediction, and add it to the current state in order to produce the predicted next state """ ### PROBLEM 1 ### YOUR CODE HERE state_ = utils.normalize(x=state, mean=self._init_dataset.state_mean, std=self._init_dataset.state_std) action_ = utils.normalize(x=action, mean=self._init_dataset.action_mean, std=self._init_dataset.action_std) input_ = tf.concat(values=[state_, action_], axis=-1) residual = utils.build_mlp(input_layer=input_, output_dim=self._state_dim, scope="dynamics", n_layers=self._nn_layers, reuse=reuse) residual = utils.unnormalize(x=residual, mean=self._init_dataset.delta_state_mean, std=self._init_dataset.delta_state_std) next_state_pred = state + residual return next_state_pred
def _dynamics_func(self, state, action, reuse): # get dataset statistics state_std = self._init_dataset.state_std state_mean = self._init_dataset.state_mean action_std = self._init_dataset.action_std action_mean = self._init_dataset.action_mean delta_state_std = self._init_dataset.delta_state_std delta_state_mean = self._init_dataset.delta_state_mean # normalize input data state_norm = utils.normalize(state, state_mean, state_std) action_norm = utils.normalize(action, action_mean, action_std) # perform delta prediction inp = tf.concat([state_norm, action_norm], 1) out = utils.build_mlp(inp, self._state_dim, 'policy', reuse=reuse) # perdict next state next_state_pred = state + utils.unnormalize(out, delta_state_mean, delta_state_std) return next_state_pred
test_dataloader = DataLoader(test_dataset, batch_size=1, shuffle=False) model.eval() result = [] for step, batch in enumerate(test_dataloader): img, mask = [t.type(torch.float).to(device) for t in batch] mask = mask.permute(0, 3, 1, 2) output, _ = model(img, mask) output_comp = mask * img + (1 - mask) * output[0] result.append(output_comp.detach().cpu()[0]) return result seed = 87 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) test_dataset = MyDataset(in_path, is_train=False) device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') model = PConvUNet().to(device) model.load_state_dict(torch.load('./model5.pth')) res = predict(model, test_dataset) for i in range(len(res)): file = os.path.join(out_path, test_dataset.imgs[i] + '.jpg') plt.imsave(file, unnormalize(res[i]))
def run_decoder_network(self, z): """ Run the decoder network to predict observations from latent states """ x = self.decoder(z) return utils.unnormalize(x, self.shift_x, self.scale_x)
def forward(self, predict, label): predict = utils.unnormalize(predict) label = utils.unnormalize(label) loss = torch.abs(predict - label) / (label + EPS) return loss.mean()
def train(self): # real_label: original img, fake_label: generated img real_label = torch.ones(self.batch_size, device=self.device) fake_label = torch.zeros(self.batch_size, device=self.device) for epoch in range(self.n_epochs): print('Epoch {} --------------------'.format(epoch + 1)) self.net_F.train() self.net_G.train() self.net_D.train() self.net_C.train() running_corrects = 0 running_loss = 0 tot_corrects = 0 tot_loss_F = 0 tot_loss_G = 0 tot_loss_D = 0 tot_loss_C = 0 tot_loss = 0 n_samples = 0 for i, (src_data, tgt_data) in enumerate(zip(self.src_train_loader, self.tgt_loader)): src_img, src_label = src_data tgt_img, _ = tgt_data # Scale img to [-1, 1] for generator src_img_scaled = (torch.stack([utils.unnormalize(img) for img in src_img.clone()]) - 0.5) * 2 src_label_onehot = torch.zeros((self.batch_size, self.n_classes + 1 )).scatter_(1, src_label.view(-1, 1), 1) tgt_label_onehot = torch.zeros((self.batch_size, self.n_classes + 1 )).index_fill_(1, torch.tensor([self.n_classes]), 1) src_img = src_img.to(self.device) src_img_scaled = src_img_scaled.to(self.device) src_label = src_label.to(self.device) tgt_img = tgt_img.to(self.device) src_label_onehot = src_label_onehot.to(self.device) tgt_label_onehot = tgt_label_onehot.to(self.device) # Update net D self.net_D.zero_grad() src_emb = self.net_F(src_img) src_emb_n_label = torch.cat((src_emb, src_label_onehot), 1) src_gen = self.net_G(src_emb_n_label) tgt_emb = self.net_F(tgt_img) tgt_emb_n_label = torch.cat((tgt_emb, tgt_label_onehot), 1) tgt_gen = self.net_G(tgt_emb_n_label) src_orig_output_D_data, src_orig_output_D_cls = self.net_D(src_img_scaled) loss_D_src_data_real = self.criterion_data(src_orig_output_D_data, real_label) loss_D_src_cls_real = self.criterion_cls(src_orig_output_D_cls, src_label) src_gen_output_D_data, src_gen_output_D_cls = self.net_D(src_gen) loss_D_src_data_fake = self.criterion_data(src_gen_output_D_data, fake_label) tgt_gen_output_D_data, _ = self.net_D(tgt_gen) loss_D_tgt_data_fake = self.criterion_data(tgt_gen_output_D_data, fake_label) loss_D = (loss_D_src_data_real + loss_D_src_cls_real + loss_D_src_data_fake + loss_D_tgt_data_fake) loss_D.backward(retain_graph=True) self.optim_D.step() # Recompute net D outputs after updating net D src_gen_output_D_data, src_gen_output_D_cls = self.net_D(src_gen) tgt_gen_output_D_data, _ = self.net_D(tgt_gen) # Update net G self.net_G.zero_grad() loss_G_data = self.criterion_data(src_gen_output_D_data, real_label) loss_G_cls = self.criterion_cls(src_gen_output_D_cls, src_label) loss_G = loss_G_data + loss_G_cls loss_G.backward(retain_graph=True) self.optim_G.step() # Update net C self.net_C.zero_grad() output_C = self.net_C(src_emb) loss_C = self.criterion_cls(output_C, src_label) loss_C.backward(retain_graph=True) self.optim_C.step() # Update net F self.net_F.zero_grad() loss_F_src = (self.adv_weight * self.criterion_cls(src_gen_output_D_cls, src_label)) loss_F_tgt = (self.adv_weight * self.alpha * self.criterion_data(tgt_gen_output_D_data, real_label)) loss_F = loss_C + loss_F_src + loss_F_tgt loss_F.backward() self.optim_F.step() # Record acc & loss _, predicts = torch.max(output_C, 1) corrects = torch.sum(predicts == src_label) running_corrects += corrects.item() tot_corrects += corrects.item() running_loss += loss_F.item() + loss_G.item() + loss_D.item() + loss_C.item() tot_loss += loss_F.item() + loss_G.item() + loss_D.item() + loss_C.item() tot_loss_F += loss_F.item() tot_loss_D += loss_D.item() tot_loss_G += loss_G.item() tot_loss_C += loss_C.item() n_samples += len(src_label) # Print record batches_to_print = max(len(self.src_train_loader), len(self.tgt_loader)) // 5 if i % batches_to_print == batches_to_print - 1: print('| {}/{} acc: {:.3f}, loss: {:.3f}' .format(i + 1, len(self.src_train_loader), running_corrects / batches_to_print / self.batch_size, running_loss / batches_to_print / self.batch_size)) running_corrects = 0 running_loss = 0 print('|__ train acc: {:.3f}, loss: {:.3f} (F: {:.3f}, D: {:.3f}, G: {:.3f}, C: {:.3f})' .format(tot_corrects / n_samples, tot_loss / n_samples, tot_loss_F / n_samples, tot_loss_D / n_samples, tot_loss_G / n_samples, tot_loss_C / n_samples)) self.hist_acc.append(tot_corrects / n_samples) self.hist_loss.append(tot_loss / n_samples) # Visualization if epoch % 5 == 0: torchvision.utils.save_image( src_img_scaled[: 64] / 2 + 0.5, self.vis_dir / 'source_e{}.png'.format(epoch + 1)) torchvision.utils.save_image( [utils.unnormalize(img.cpu()) for img in tgt_img[: 64]], self.vis_dir / 'target_e{}.png'.format(epoch + 1)) torchvision.utils.save_image( src_gen[: 64] / 2 + 0.5, self.vis_dir / 'source_gen_e{}.png'.format(epoch + 1)) torchvision.utils.save_image( tgt_gen[: 64] / 2 + 0.5, self.vis_dir / 'target_gen_e{}.png'.format(epoch + 1)) # Validate per epoch self.validate(epoch + 1, tot_loss / n_samples) # Process Visualization import matplotlib.pyplot as plt x_range = range(1, epoch + 2) plt.plot(x_range, self.hist_acc) plt.savefig(str(self.vis_dir / 'train_acc')) plt.close() plt.plot(x_range, self.hist_loss) plt.savefig(str(self.vis_dir / 'train_loss')) plt.close()
# Hyperparameters learning_rate = 0.5 generation = 200 # Learning loop cost_history = [] accuracy_history = [] curve_history = [] for i in range(generation): sum_for_t0 = 0 sum_for_t1 = 0 current_curve = [] for km, price in zip(n_km_list, n_price_list): prediction = predict(km, t0, t1) current_curve.append( unnormalize(prediction, price_decay, price_max)) error = prediction - price sum_for_t0 += error sum_for_t1 += error * km # Regression t0 -= learning_rate * (data_mean * sum_for_t0) t1 -= learning_rate * (data_mean * sum_for_t1) # Evolution debug sum_for_t0 *= data_mean precision = 1 - abs(sum_for_t0**2) print( "generation {}: error = {:.2f}, precision = {:.2f}, T0 = {:.2f}, T1 = {:.2f}" .format(i, sum_for_t0, precision, t0, t1))