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))
