def prepare(args):
dataset_path = args.dataset_path # '../ml-100k'
save_path = args.save_processed_data_path # './weights'
user_item_matrix, raw_side_feature_v, raw_side_feature_u = preprocess(dataset_path, save_path)
# user_item_matrix[943, 1682], raw_side_feature_u[943, 21], raw_side_feature_u[1682,19]
num_user, num_item = user_item_matrix.shape
mask = user_item_matrix > 0
mask_new = mask + np.random.uniform(0, 1, (num_user, num_item))
train_mask = (mask_new <= (1 + args.split_ratio)) & mask
test_mask = (mask_new > (1 + args.split_ratio)) & mask
user_item_matrix_train = user_item_matrix + 0
user_item_matrix_train[test_mask] = 0
user_item_matrix_test = user_item_matrix + 0
user_item_matrix_test[train_mask] = 0
all_M_u = []
all_M_v = []
all_M = []
for i in range(args.rate_num): # rate_num==5, 5分制
M_r = user_item_matrix_train == (i + 1) # 用户u给物品v评分为r
all_M_u.append(utils.normalize(M_r)) # 评分
all_M_v.append(utils.normalize(M_r.T)) # 被评分
all_M.append(M_r)
# [[943,1682],[943,1682],[943,1682],[943,1682],[943,1682]] -> ndarray(5, 943, 1682), 5分制
all_M = np.array(all_M)
mask = user_item_matrix_train > 0
## 读取side_information并处理
side_feature_u = raw_side_feature_u
side_feature_v = raw_side_feature_v
adjacency_u = epsilon_similarity_graph(side_feature_u, epsilon=1.1)
laplacian_u = compute_laplacian(adjacency_u, True)
adjacency_v = epsilon_similarity_graph(side_feature_v, epsilon=2.1)
laplacian_v = compute_laplacian(adjacency_v, True)
laplacian_u = utils.np_to_var(laplacian_u)
laplacian_v = utils.np_to_var(laplacian_v)
## 产生输入特征
# 对应论文中的(Xu;Xv) = I
feature_dim = num_user + num_item
I = np.eye(num_user + num_item)
feature_u = I[0:num_user, :]
feature_v = I[num_user:, :]
return feature_u, feature_v, feature_dim, all_M_u, all_M_v, side_feature_u, side_feature_v, all_M, mask, user_item_matrix_train, user_item_matrix_test, laplacian_u, laplacian_v
def rollout(env,
start_state: State,
policy,
buffer,
num_steps=args.subpolicy_duration) -> None:
s = start_state
for i in range(num_steps):
mean, std = policy(np_to_var(s))
p = distributions.Normal(mean, std)
a = p.sample()
succ, r, done, _ = env.step(a.data.numpy())
env.done = done
buffer.append(Step(s, a, r, succ, done))
s = np_to_var(env.reset()) if done else succ
def test_batch_with_labels(net, file, resize = False, batch_size = 10, image_size = 384, smooth = 1.0, lam = 1.0):
'''
Test on a validation dataset (here we only consider BCE loss instead of focal loss).
No TTA or ensemble used at this case.
Parameters:
@net: the object for network model.
@file: root directory of the validation dataset.
@resize: boolean flag for image resize.
@batch_size: batch size
@image_size: the size that the image is converted to.
@smooth: number to be added on denominator and numerator when compute dice loss.
@lam: weight to balance the dice loss in the final combined loss.
Return:
average loss (BCE + dice) over batches
F1 score of the test
'''
# keep original data
data_augment = False
rotate = False
change_color = False
test_dataset = utils.MyDataset(file, resize, data_augment, image_size, rotate, change_color)
dataloader = utils_data.DataLoader(dataset = test_dataset, batch_size = batch_size, shuffle=False)
epoch_loss = 0.0
numer = 0.0
denom = 0.0
gamma = 0.0
loss_type = 'bce'
Loss = utils.loss(smooth, lam, gamma, loss_type)
for i, batch in enumerate(dataloader):
print('Test on batch %d'%i)
image = utils.np_to_var(batch['image'])
mask = utils.np_to_var(batch['mask'])
pred = net.forward(image)
loss = Loss.final_loss(pred, mask)
epoch_loss += loss.data.item() * batch_size
mask = utils.var_to_np(mask)
pred = utils.var_to_np(pred)
numer += (mask * (pred > 0.5)).sum()
denom += mask.sum() + (pred > 0.5).sum()
epoch_loss /= len(test_dataset)
f1 = 2.0 * numer / denom
return epoch_loss, f1
def evaluate_model(net, net_input, img_np, img_noisy_np, num_iter=6000,
show_every=500, report=True, figsize=10):
loss_fn=torch.nn.MSELoss().type(dtype)
input_noise=True
LR=0.01
reg_noise_std=1./30.
net_input_saved = net_input.data.clone()
noise = net_input.data.clone()
img_noisy_var = utils.np_to_var(img_noisy_np).type(dtype)
psnr_history = []
def closure(i):
if input_noise:
net_input.data = net_input_saved + (noise.normal_() * reg_noise_std)
out = net(net_input)
total_loss = loss_fn(out, img_noisy_var)
total_loss.backward()
psrn = compare_psnr(utils.var_to_np(out), img_np)
psnr_history.append(psrn)
if report:
print ('Iteration %05d Loss %f PSNR %.3f' % (i, total_loss.data[0], psrn), '\r', end='')
if i % show_every == 0:
out_np = utils.var_to_np(out)
utils.plot_image_grid([np.clip(out_np, 0, 1)], factor=figsize, nrow=1)
return total_loss
print('Starting optimization with ADAM')
optimizer = torch.optim.Adam(net.parameters(), lr=LR)
for j in range(num_iter):
optimizer.zero_grad()
closure(j)
optimizer.step()
if report:
out_np = utils.var_to_np(net(net_input))
q = utils.plot_image_grid([np.clip(out_np, 0, 1), img_np], factor=13);
data = {}
data['psnr_history'] = psnr_history
pickle.dump(data, open('denoising_psnr.p','wb'))
max_index, max_value = max(enumerate(psnr_history), key=operator.itemgetter(1))
return max_index, max_value
def test_single_image(net, file, size = 384, resize = True):
'''
Test a single image with the trained model.
Parameters:
@net: the object for network model.
@file: name of the image file for test.
@size: the size that the image is converted to.
@resize: boolean flag for image resize.
Return:
predicted segmentation mask
original image covered by the red mask
'''
uint_image = io.imread(file)
test_image_origin = np.array(uint_image).astype(np.float32) / 255.0
# convert the resized image to a Torch tensor with batch size 1
if resize:
test_image_origin = transform.resize(test_image_origin, (size, size), mode = 'constant', anti_aliasing = True)
test_image = np.moveaxis(test_image_origin, 2, 0).astype(np.float32) # tensor format
else:
test_image = test_image_origin
test_image = np.moveaxis(test_image, 2, 0).astype(np.float32) # tensor format
test_image = np.expand_dims(test_image, axis = 0)
test_image = utils.np_to_var(torch.from_numpy(test_image))
pred_test = net.forward(test_image)
pred_np = utils.var_to_np(pred_test)[0][0] # predicted confidence map
# convert prediction to binary segmentation mask
new_mask = (pred_np >= 0.5)
# make a red cover on the original image for visualization of segmentation result
channel = test_image_origin[:, :, 0]
channel[new_mask] = 1.0
test_image_origin[:, :, 0] = channel
mask = new_mask * 255
return mask.astype(np.uint8), test_image_origin
s = np_to_var(env.reset()) if done else succ
if __name__ == "__main__":
envs = [gym.make(env) for env in args.env_names]
# TODO assert that all envs have same state and action spaces
# Start with random env
for env in envs:
env.done = False
env = random.choice(envs)
S = int(np.prod(env.observation_space.shape))
A = int(np.prod(env.action_space.shape))
H = hidden_size = 50
policies = [Policy(S, H, A) for _ in range(args.num_policies)]
buffer = ReplayBuffer(args.max_buffer_size)
s = env.reset()
master = MasterPolicy(S, H, args.num_policies)
P = master(np_to_var(s))
for t in itertools.count():
# pick new subpolicy every X steps
if t % args.subpolicy_duration == 0:
policy = policies[int(P.sample())]
rollout(env, s, policy, buffer, args.subpolicy_duration)
def prepare(args):
dataset_path = args.dataset_path
save_path = args.save_processed_data_path
user_item_matrix, raw_side_feature_v, raw_side_feature_u = preprocess(
dataset_path, save_path)
num_user, num_item = user_item_matrix.shape
mask = user_item_matrix > 0
mask_new = mask + np.random.uniform(0, 1, (num_user, num_item))
train_mask = (mask_new <= (1 + args.split_ratio)) & mask
test_mask = (mask_new > (1 + args.split_ratio)) & mask
user_item_matrix_train = user_item_matrix + 0
user_item_matrix_train[test_mask] = 0
user_item_matrix_test = user_item_matrix + 0
user_item_matrix_test[train_mask] = 0
all_M_u = []
all_M_v = []
all_M = []
for i in range(args.rate_num):
M_r = user_item_matrix_train == (i + 1)
all_M_u.append(utils.normalize(M_r))
all_M_v.append(utils.normalize(M_r.T))
all_M.append(M_r)
all_M = np.array(all_M)
mask = user_item_matrix_train > 0
### side feature loading and processing
if args.use_data_whitening:
print("Using data whitening!")
M_u, mean_u = data_whitening(raw_side_feature_u)
M_v, mean_v = data_whitening(raw_side_feature_v)
side_feature_u = np.dot(raw_side_feature_u - mean_u, M_u)
side_feature_v = np.dot(raw_side_feature_v - mean_v, M_v)
else:
print("Not using data whitening!")
side_feature_u = raw_side_feature_u
side_feature_v = raw_side_feature_v
############test############
if args.use_data_whitening:
adjacency_u = epsilon_similarity_graph(side_feature_u, epsilon=5.7)
laplacian_u = compute_laplacian(adjacency_u, True)
adjacency_v = epsilon_similarity_graph(side_feature_v, epsilon=52.5)
laplacian_v = compute_laplacian(adjacency_v, True)
else:
adjacency_u = epsilon_similarity_graph(side_feature_u, epsilon=1.1)
laplacian_u = compute_laplacian(adjacency_u, True)
adjacency_v = epsilon_similarity_graph(side_feature_v, epsilon=2.1)
laplacian_v = compute_laplacian(adjacency_v, True)
laplacian_u = utils.np_to_var(laplacian_u)
laplacian_v = utils.np_to_var(laplacian_v)
### input feature generation
feature_dim = num_user + num_item
I = np.eye(num_user + num_item)
feature_u = I[0:num_user, :]
feature_v = I[num_user:, :]
return feature_u, feature_v, feature_dim, all_M_u, all_M_v, side_feature_u, side_feature_v, all_M, mask, user_item_matrix_train, user_item_matrix_test, laplacian_u, laplacian_v
def test_single_with_TTA(net, file, size = 384, resize = True):
'''
Test a single image with the trained model, using test time augmentation (TTA).
Parameters:
@net: the object for network model.
@file: name of the image file for test.
@size: the size that the image is converted to.
@resize: boolean flag for image resize.
Return:
predicted segmentation mask
original image covered by the red mask
'''
uint_image = io.imread(file)
test_image_origin = np.array(uint_image).astype(np.float32) / 255.0
# resize
if resize:
test_image = transform.resize(test_image_origin, (size, size), mode = 'constant', anti_aliasing = True)
else:
test_image = test_image_origin
# create tensor for collecting 8 images (generated by TTA)
image_set = []
for i in range(8):
b1 = i // 4
b2 = (i - b1 * 4) // 2
b3 = i - b1 * 4 - b2 * 2
# flip and rotate the image based on perturbed choice
tem_image = utils.flip_rotate(test_image, b1, b2, b3, inverse = False)
# convert to tensor format
tem_image = np.moveaxis(tem_image, 2, 0).astype(np.float32)
image_set.append(tem_image)
image_tensor = np.array(image_set)
# convert to Torch tensor variable
image_tensor = utils.np_to_var(torch.from_numpy(image_tensor))
pred_test = net.forward(image_tensor)
pred_np = utils.var_to_np(pred_test)
pred = np.squeeze(pred_np, axis = 1)
# inversely transform each prediction back to the original orientation
for i in range(8):
b1 = i // 4
b2 = (i - b1 * 4) // 2
b3 = i - b1 * 4 - b2 * 2
pred[i] = utils.flip_rotate(pred[i], b1, b2, b3, inverse = True)
# merge into one prediction
pred = np.median(pred, axis = 0)
# convert prediction to binary segmentation mask
new_mask = (pred >= 0.5)
# make a red cover on the original image for visualization of segmentation result
channel = test_image_origin[:, :, 0]
channel[new_mask] = 1.0
test_image_origin[:, :, 0] = channel
mask = new_mask * 255
return mask.astype(np.uint8), test_image_origin
def evaluate_model(net, net_input, img_np, img_mask_np, num_iter=6000,
show_every=500, report=True, figsize=10):
pad = 'zero'
OPTIMIZER = 'adam'
INPUT = 'noise'
input_depth = 32
LR = 0.01
num_iter = 6001
param_noise = False
show_every = 500
figsize = 10
net_input_saved = net_input.data.clone()
noise = net_input.data.clone()
#img_noisy_var = utils.np_to_var(img_noisy_np).type(dtype)
# img
psnr_history = []
# Loss
mse = torch.nn.MSELoss().type(dtype)
img_var = utils.np_to_var(img_np).type(dtype)
mask_var = utils.np_to_var(img_mask_np).type(dtype)
psnr_history = []
def closure(i):
#global i
if param_noise:
for n in [x for x in net.parameters() if len(x.size()) == 4]:
n.data += n.data.clone().normal_()*n.data.std()/50
out = net(net_input)
total_loss = mse(out * mask_var, img_var * mask_var)
total_loss.backward()
psrn = compare_psnr(utils.var_to_np(out), img_np)
psnr_history.append(psrn)
if report:
print ('Iteration %05d Loss %f PSNR %.3f' % (i, total_loss.data[0], psrn), '\r', end='')
if i % show_every == 0:
out_np = utils.var_to_np(out)
utils.plot_image_grid([np.clip(out_np, 0, 1)], factor=figsize, nrow=1)
#i += 1
return total_loss
print('Starting optimization with ADAM')
optimizer = torch.optim.Adam(net.parameters(), lr=LR)
for j in range(num_iter):
optimizer.zero_grad()
closure(j)
optimizer.step()
if report:
out_np = utils.var_to_np(net(net_input))
q = utils.plot_image_grid([np.clip(out_np, 0, 1), img_np], factor=13);
data = {}
data['psnr_history'] = psnr_history
pickle.dump(data, open('inpainting_validation_psnr.p','wb'))
max_index, max_value = max(enumerate(psnr_history), key=operator.itemgetter(1))
return max_index, max_value
def train_net(root,
resize,
data_augment,
rotate,
change_color,
lr,
weight_decay,
model_choice,
save_ckpt,
image_size,
batch_size,
num_epochs,
save_test_image,
test_image_name,
early_stop,
early_stop_tol,
lr_decay,
decay_rate,
decay_period,
validate_root,
loss_type='bce',
smooth=1.0,
lam=1.0,
gamma=2.0):
'''
Network training, which will output:
1. log for loss in every iteration, in text file.
2. saved checkpoint which contains the trained parameters, in directory ./parameters
3. segmentation result on the test image, saved in directory ./epoch_output.
Parameters:
@root: root directory for training dataset.
@resize: boolean flag for image resizing.
@data_augment: boolean flag for DA8 (randomly rotate 90 degrees, flip horizontally and vertically).
@rotate: boolean flag for random rotation to the training images.
@change_color: boolean flag for random perturbation on HSV channels of the training images.
@lr: learning rate.
@weight_decay: weight decay for L2 regularization on the network parameters.
@model_choice: 1 for LinkNet, 2 for D-LinkNet, 3 for D-LinkNet+.
@save_ckpt: the period (in epochs) to save the checkpoint of the network.
@image_size: the image size for the images to trained.
@batch_size: batch size for mini-batch stochastic gradient descent.
@num_epochs: number of epochs for training.
@save_test_image: the period (in epochs) to save the prediction of the test image.
@test_image_name: the name of the test image.
@early_stop: the boolean flag to have early stop.
@early_stop_tol: the tolerance (in number of saving checkpoints) to trigger early stop.
@lr_decay: boolean flag for learning rate decay in every decay period.
@decay_rate: decay ratio for learning rate, e.g. lr = lr * lr_decay.
@decay_period: the period in number of epochs to trigger the learning rate decay.
@validate_root: root directory for validation dataset (mainly for evaluation of network during training).
@loss_type: either 'bce' (BCE loss) or 'focal' (focal loss).
@smooth: number to be added on denominator and numerator when compute dice loss.
@lam: weight to balance the dice loss in the final combined loss.
@gamma: for focal loss.
'''
if os.path.exists('./epoch_output'):
shutil.rmtree('./epoch_output')
os.makedirs('./epoch_output')
if not os.path.exists('./parameters'):
os.makedirs('./parameters')
weights_name = './parameters/weights' + str(model_choice)
net = utils.create_models(model_choice)
net.train() # in train mode
# net = torch.nn.DataParallel(net, device_ids=range(torch.cuda.device_count()))
# create AMSGrad optimizer
optimizer = optim.Adam(net.parameters(),
lr=lr,
weight_decay=weight_decay,
amsgrad=True)
Loss = utils.loss(smooth, lam, gamma, loss_type)
dataloader = utils.get_data_loader(root, resize, data_augment, image_size,
batch_size, rotate, change_color)
num_batch = len(dataloader)
total_train_iters = num_epochs * num_batch
loss_history = []
print('Started training at {}'.format(
time.asctime(time.localtime(time.time()))))
test_loss = 100.0
count = 0
for epoch in range(num_epochs):
print('Start epoch ', epoch)
epoch_loss = 0
t = time.time()
for iteration, batch in enumerate(dataloader, epoch * num_batch + 1):
print('Iteration: ', iteration)
print('Time for loading the data takes: ', time.time() - t, ' s')
t = time.time()
image = utils.np_to_var(batch['image'])
mask = utils.np_to_var(batch['mask'])
optimizer.zero_grad()
pred = net.forward(image)
loss = Loss.final_loss(pred, mask)
loss.backward()
optimizer.step()
epoch_loss += loss.data.item()
# print the log info
print('Iteration [{:6d}/{:6d}] | loss: {:.4f}'.format(
iteration, total_train_iters, loss.data.item()))
print('Time spent on back propagation: ', time.time() - t, ' s')
loss_history.append(loss.data.item())
t = time.time()
# save the test image for visualizing the training outcome
if (epoch + 1) % save_test_image == 0:
with torch.no_grad():
_, test_image = test.test_single_image(net,
test_image_name,
resize=False)
io.imsave('./epoch_output/test_epoch' + str(epoch) + '.png',
test_image)
# early stop
if early_stop and (epoch + 1) % save_ckpt == 0:
with torch.no_grad():
loss, f1 = test.test_batch_with_labels(net,
validate_root,
resize=False,
batch_size=10,
image_size=image_size,
smooth=smooth,
lam=lam)
print('On the validation dataset, loss: ', loss, ', F1: ', f1)
if loss <= test_loss:
test_loss = loss
count = 0
torch.save(net.state_dict(), weights_name)
elif count < early_stop_tol:
count += 1
lr *= decay_rate
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('The new loss is found to be larger than before')
else:
print('Reach the early stop tolerence...')
print('Break the update at ', epoch, 'th epoch')
break
if not early_stop and (epoch + 1) % save_ckpt == 0:
with torch.no_grad():
torch.save(net.state_dict(), weights_name)
if lr_decay and (epoch + 1) % decay_period == 0:
with torch.no_grad():
lr *= decay_rate
for param_group in optimizer.param_groups:
param_group['lr'] = lr
epoch_loss /= num_batch
print('In the epoch ', epoch, ', the average batch loss is ',
epoch_loss)
if not early_stop:
torch.save(net.state_dict(), weights_name)
# save the loss history
with open('loss.txt', 'wt') as file:
file.write('\n'.join(['{}'.format(loss) for loss in loss_history]))
file.write('\n')
return states, actions, rewards, succ_states, dones
def get_critic_train_data(succ_states, rewards, dones):
# r + Q(s, pi(s'))
Q_succ = critic_target(succ_states, actor_target(succ_states)).squeeze()
td_estimate = rewards + ((1 - dones) * hparams.discount * Q_succ)
return td_estimate.detach()
actor_opt = Adam(actor.parameters())
critic_opt = Adam(critic.parameters())
buffer = ReplayBuffer(hparams.buffer_size)
s, rews = np_to_var(env.reset()), []
for hparam in hparams.trials(5):
exp.add_argparse_meta(hparam)
for timestep in range(hparam.num_steps):
noise = Normal(
mean=Variable(torch.zeros(A)),
std=hparam.noise_factor * Variable(torch.ones(A)),
)
if timestep % 1000 == 0:
hparam.noise_factor /= 2
a = actor(s) + noise.sample()
succ, r, done, _ = env.step(a.data.numpy())
succ = np_to_var(succ)