def main():
nclass = 7
batch_size = 4
drop_last = True
train_loader, val_loader = make_dataloader(batch_size=batch_size,
drop_last=drop_last)
model = DeepLab(num_classes=nclass, output_stride=16, freeze_bn=False)
model = load_pretrained_model(model)
evaluator = Evaluator(nclass)
model = model.cuda()
model.eval()
evaluator.reset()
for i, sample in enumerate(val_loader):
image, target = sample['image'], sample['label']
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = model(image)
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
evaluator.add_batch(target, pred)
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
def __init__(self, args):
self.args = args
# define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args)
# define network
model = DeepLab(num_classes=self.nclass, output_stride=args.out_stride)
optimizer = torch.optim.Adam(model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay)
self.criterion = SegmentationLoss(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
self.evaluator = Evaluator(self.nclass)
self.best_pred = 0.0
self.trainloss_history = []
self.valloss_history = []
self.train_plot = []
self.val_plot = []
# every 10 epochs the lr will multiply 0.1
self.scheduler = LR_Scheduler(args.lr_scheduler, args.base_lr, args.epochs, len(self.train_loader), lr_step=20)
if args.cuda:
self.model = self.model.cuda()
def __init__(self, env_cls, args, manager, cfg, process_num, pre=False, pre_irl=False, infer=False, realtrain=False):
"""
:param env_cls: env class or function, not instance, as we need to create several instance in class.
:param args:
:param manager:
:param cfg:
:param process_num: process number
:param pre: set to pretrain mode
:param infer: set to test mode
"""
self.train_direc = args.train_direc
self.realtrain = realtrain
self.process_num = process_num
self.human_reward = human_reward(args, cfg)
# initialize envs for each process
self.env_list = []
for _ in range(process_num):
self.env_list.append(env_cls())
# construct policy and value network
self.feature_layer = nn.Sequential(
nn.Linear(cfg.s_dim, cfg.h_dim),
nn.ReLU()
)
self.policy = MultiDiscretePolicy(cfg, self.feature_layer).to(device=DEVICE)
self.value = Value(cfg).to(device=DEVICE)
self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
if pre:
self.print_per_batch = args.print_per_batch
from dbquery import DBQuery
db = DBQuery(args.data_dir)
self.data_train = manager.create_dataset_rl('train', args.batchsz, cfg, db)
self.data_valid = manager.create_dataset_rl('valid', args.batchsz, cfg, db)
self.data_test = manager.create_dataset_rl('test', args.batchsz, cfg, db)
# self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
else:
self.rewarder = RewardEstimator(args, manager, cfg, pretrain=pre_irl, inference=infer, feature_extractor=self.policy)
# self.rewarder = DiscEstimator(args, manager, cfg, pretrain=pre_irl, inference=infer)
self.evaluator = Evaluator(args.data_dir, cfg)
self.save_dir = args.save_dir
self.save_per_epoch = args.save_per_epoch
self.optim_batchsz = args.batchsz
self.update_round = args.update_round
self.policy.eval()
self.value.eval()
self.gamma = args.gamma
self.epsilon = args.epsilon
self.tau = args.tau
self.policy_optim = optim.RMSprop(self.policy.parameters(), lr=args.lr_rl)
self.value_optim = optim.Adam(self.value.parameters(), lr=args.lr_rl)
self.mt_factor = args.mt_factor
def evaluate_batch(gt_list,pred_list,num_of_class):
evaluator = Evaluator(num_of_class)
evaluator.reset()
for i in range(len(gt_list)):
evaluator.add_batch(gt_list[i],pred_list[i])
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
return format(Acc, '.4f'), format(Acc_class, '.4f'), format(mIoU, '.4f'), format(FWIoU, '.4f')
def evaluate_single(gt,pred,num_of_class):
evaluator = Evaluator(num_of_class)
evaluator.reset()
evaluator.add_batch(gt,pred)
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
return format(Acc, '.4f'), format(Acc_class, '.4f'), format(mIoU, '.4f'), format(FWIoU, '.4f')
def create_evaluator():
return Evaluator(
graph.n_node,
given_as_target=given_as_target,
siblings_nodes=siblings_nodes,
config=config,
)
def __init__(self,
anchors,
size: Tuple[int, int],
metric_names: list,
detect_thresh: float = 0.3,
nms_thresh: float = 0.3,
images_per_batch: int = -1):
self.ap = 'AP'
self.anchors = anchors
self.size = size
self.detect_thresh = detect_thresh
self.nms_thresh = nms_thresh
self.images_per_batch = images_per_batch
self.metric_names_original = metric_names
self.metric_names = ["{}-{}".format(self.ap, i) for i in metric_names]
self.evaluator = Evaluator()
self.boundingBoxes = BoundingBoxes()
def __init__(self, args):
self.epochs = args.epochs
self.save_interval = args.save_interval
self.train_data = AerialDataset(args, mode='train')
self.train_loader = DataLoader(self.train_data,
batch_size=args.train_batch_size,
shuffle=True,
num_workers=1)
self.model = models.deeplabv3_resnet50(num_classes=args.num_of_class)
#self.model = models.fcn_resnet50(num_classes=args.num_of_class)
self.loss = args.loss
if self.loss == 'CE':
self.criterion = nn.CrossEntropyLoss()
else: #self.loss == 'LS'
self.criterion = LovaszSoftmax()
self.optimizer = torch.optim.AdamW(self.model.parameters())
self.eval_interval = args.eval_interval
self.eval_data = AerialDataset(args, mode='eval')
self.eval_loader = DataLoader(self.eval_data,
batch_size=args.eval_batch_size,
shuffle=False,
num_workers=1)
self.evaluator = Evaluator(args.num_of_class)
self.cuda = args.cuda
if self.cuda is True:
self.model = self.model.cuda()
self.resume = args.resume
if self.resume != None:
if self.cuda:
checkpoint = torch.load(args.resume)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
self.model.load_state_dict(checkpoint['model'])
self.optimizer.load_state_dict(checkpoint['opt'])
self.start_epoch = checkpoint['epoch'] + 1
#start from next epoch
else:
self.start_epoch = 1
class human_reward(object):
def __init__(self, args, config):
self.evaluator = Evaluator(args.data_dir, config)
def check_success(self, s):
match_session = self.evaluator.match_rate(s, True)
inform_session = self.evaluator.inform_F1(s, True)
if (match_session == 1 and inform_session[1] == 1) \
or (match_session == 1 and inform_session[1] is None) \
or (match_session is None and inform_session[1] == 1):
return True
else:
return False
def reward_human(self, s, done):
success = self.check_success(s)
if success:
reward = 0.2 * 40
if not success and not done:
reward = -0.1
if not success and done:
reward = -0.1 * 40 # 10
return reward
def evaluate(image_file_paths,
gt_file_paths,
model,
nn_input_size,
num_classes,
window_sizes=None,
step_sizes=None,
device=None):
"""
Evaluate model performance
:param image_file_paths: paths to images that will be predicted
:param gt_file_paths: paths to ground truth masks
:param model: model used for prediction
:param nn_input_size: size that the images will be scaled to before feeding them into the nn
:param num_classes: number of classes
:param window_sizes: list of sizes that determine how the image will be cut (for sliding window). Image will be cut
into squares
:param step_sizes: list of step sizes for sliding window
:param device: PyTorch device (cpu or gpu)
:return: list of prediction masks if res_fcn is None, else Nothing
"""
model.eval()
evaluator = Evaluator(num_classes)
def ev(prediction, index):
pred_class_id_mask = np.argmax(prediction, axis=-1)
gt_colour_mask = cv2.imread(gt_file_paths[index])[:, :, 0]
gt_class_id_mask = colour_mask_to_class_id_mask(gt_colour_mask)
# Add sample into evaluator
evaluator.add_batch(gt_class_id_mask, pred_class_id_mask)
with torch.no_grad():
predict.predict(image_file_paths=image_file_paths,
model=model,
nn_input_size=nn_input_size,
res_fcn=ev,
window_sizes=window_sizes,
step_sizes=step_sizes,
device=device)
with np.errstate(divide='ignore', invalid='ignore'):
acc = evaluator.Pixel_Accuracy()
acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
fWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
return acc, acc_class, mIoU, fWIoU
def func_wrapper(model, arguments):
evaluator = Evaluator(args.num_classes)
evaluator.reset()
model.eval()
model.cuda()
threshold, use_cuda = arguments[0], arguments[1]
total_samples = 0
for sample in tqdm(data_loader):
images, label = sample['image'], sample['label']
images, label = images.cuda(), label.cpu().numpy()
output = model(images)
pred = torch.argmax(output, 1).data.cpu().numpy()
evaluator.add_batch(label, pred)
total_samples += images.size()[0]
if total_samples > threshold:
break
mIoU = evaluator.Mean_Intersection_over_Union() * 100.
print("mIoU : {:0.2f}".format(mIoU))
return mIoU
def main():
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
testdataset = Cityscapesloader('Cityscapes/', split='val', )
testdataloader = DataLoader(testdataset, batch_size=4, shuffle=True, num_workers=8, pin_memory=True)
ICnet_ = ICnet_model(19, (512, 1024)).cuda()
model_path = 'checkpoints/last.pth'
if os.path.exists(model_path):
ICnet_.load(model_path)
summary(ICnet_, ( 3, 512, 1024))
ICnet_.eval()
eval = Evaluator(19)
eval.reset()
with torch.no_grad():
for batch_index, (image_batch, labels_batch) in enumerate(testdataloader):
print(batch_index)
image = image_batch.cuda()
labels_batch = np.array(labels_batch)
output = ICnet_.forward(image).cpu().numpy()
pred = np.argmax(output, axis=1)
eval.add_batch(labels_batch, pred)
Pixel_Accuracy = eval.Pixel_Accuracy()
print('Pixel_Accuracy = ', Pixel_Accuracy)
def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
load_weights_folder = os.path.expanduser(settings.load_weights_dir)
model_path = os.path.join(load_weights_folder, "model.pth")
model_dict = torch.load(model_path)
# data
datasets_dict = {
"3d60": datasets.ThreeD60,
"panosuncg": datasets.PanoSunCG,
"stanford2d3d": datasets.Stanford2D3D,
"matterport3d": datasets.Matterport3D
}
dataset = datasets_dict[settings.dataset]
fpath = os.path.join(os.path.dirname(__file__), "datasets", "{}_{}.txt")
test_file_list = fpath.format(settings.dataset, "test")
test_dataset = dataset(settings.data_path,
test_file_list,
model_dict['height'],
model_dict['width'],
is_training=False)
test_loader = DataLoader(test_dataset,
settings.batch_size,
False,
num_workers=settings.num_workers,
pin_memory=True,
drop_last=False)
num_test_samples = len(test_dataset)
num_steps = num_test_samples // settings.batch_size
print("Num. of test samples:", num_test_samples, "Num. of steps:",
num_steps, "\n")
# network
Net_dict = {"UniFuse": UniFuse, "Equi": Equi}
Net = Net_dict[model_dict['net']]
model = Net(model_dict['layers'],
model_dict['height'],
model_dict['width'],
max_depth=test_dataset.max_depth_meters,
fusion_type=model_dict['fusion'],
se_in_fusion=model_dict['se_in_fusion'])
model.to(device)
model_state_dict = model.state_dict()
model.load_state_dict(
{k: v
for k, v in model_dict.items() if k in model_state_dict})
model.eval()
evaluator = Evaluator(settings.median_align)
evaluator.reset_eval_metrics()
saver = Saver(load_weights_folder)
pbar = tqdm.tqdm(test_loader)
pbar.set_description("Testing")
with torch.no_grad():
for batch_idx, inputs in enumerate(pbar):
equi_inputs = inputs["normalized_rgb"].to(device)
cube_inputs = inputs["normalized_cube_rgb"].to(device)
outputs = model(equi_inputs, cube_inputs)
pred_depth = outputs["pred_depth"].detach().cpu()
gt_depth = inputs["gt_depth"]
mask = inputs["val_mask"]
for i in range(gt_depth.shape[0]):
evaluator.compute_eval_metrics(gt_depth[i:i + 1],
pred_depth[i:i + 1],
mask[i:i + 1])
if settings.save_samples:
saver.save_samples(inputs["rgb"], gt_depth, pred_depth, mask)
evaluator.print(load_weights_folder)
class Trainer(object):
def __init__(self, args):
self.args = args
# define Dataloader
self.train_loader, self.val_loader, self.test_loader, self.nclass = make_data_loader(args)
# define network
model = DeepLab(num_classes=self.nclass, output_stride=args.out_stride)
optimizer = torch.optim.Adam(model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay)
self.criterion = SegmentationLoss(weight=None, cuda=args.cuda).build_loss(mode=args.loss_type)
self.model, self.optimizer = model, optimizer
self.evaluator = Evaluator(self.nclass)
self.best_pred = 0.0
self.trainloss_history = []
self.valloss_history = []
self.train_plot = []
self.val_plot = []
# every 10 epochs the lr will multiply 0.1
self.scheduler = LR_Scheduler(args.lr_scheduler, args.base_lr, args.epochs, len(self.train_loader), lr_step=20)
if args.cuda:
self.model = self.model.cuda()
def training(self, epoch):
train_loss = 0.0
self.model.train()
#tbar = tqdm(self.train_loader) # Instantly make your loops show a smart progress meter
for i, sample in enumerate(self.train_loader):
if i >= 10000:
break
else:
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
self.scheduler(self.optimizer, i, epoch, self.best_pred)
self.optimizer.zero_grad()
output = self.model(image)
loss = self.criterion(output, target)
loss.backward()
self.optimizer.step()
self.trainloss_history.append(loss.data.cpu().numpy())
if epoch == 3 and i == 1:
print("after 3 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 9 and i == 1:
print("after 10 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 29 and i == 1:
print("after 20 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 29 and i == 3:
visualize_image("coco", image, target, output)
if epoch == 29 and i == 5:
visualize_image("coco", image, target, output)
if epoch == 29 and i == 7:
visualize_image("coco", image, target, output)
if epoch == 29 and i == 9:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 1:
print("after 30 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 49 and i == 3:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 5:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 7:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 9:
visualize_image("coco", image, target, output)
pred = output.data.cpu().numpy()
pred = np.argmax(pred, axis=1)
self.evaluator.add_batch(target.cpu().numpy(), pred)
acc = self.evaluator.Pixel_Accuracy()
#acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
last_loss = self.trainloss_history[-i:]
train_loss = np.mean(last_loss)
self.train_plot.append(train_loss)
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size ))
print('Loss: %.3f' % train_loss)
print("Train_Acc:{}".format(acc), "Train_mIoU:{}".format(mIoU))
if epoch == 20:
torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_20')
if epoch == 30:
torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_30')
if epoch == 40:
torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_40')
if epoch == 50:
torch.save(self.model.state_dict(), '/home/wan/Segmentation/model_50')
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
#tbar = tqdm(self.val_loader, desc='\r')
val_loss = 0.0
for i, sample in enumerate(self.val_loader):
if i >= 750:
break
else:
image, target = sample['image'], sample['label']
if self.args.cuda:
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = self.model(image)
loss = self.criterion(output, target)
self.valloss_history.append(loss.data.cpu().numpy())
pred = output.data.cpu().numpy()
pred = np.argmax(pred, axis=1)
# visualize the prediction and target
if epoch == 3 and i == 0:
print("after 3 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch ==9 and i == 0:
print("after 10 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 29 and i == 0:
print("after 20 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 29 and i == 1:
visualize_image("coco", image, target, output)
if epoch == 29 and i == 2:
visualize_image("coco", image, target, output)
if epoch == 29 and i == 3:
visualize_image("coco", image, target, output)
if epoch == 29 and i == 4:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 0:
print("after 30 epochs, the result of training data:")
visualize_image("coco", image, target, output)
if epoch == 49 and i == 1:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 2:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 3:
visualize_image("coco", image, target, output)
if epoch == 49 and i == 4:
visualize_image("coco", image, target, output)
self.evaluator.add_batch(target.cpu().numpy(), pred)
acc = self.evaluator.Pixel_Accuracy()
#acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
last_loss = self.valloss_history[-i:]
val_loss = np.mean(last_loss)
self.val_plot.append(val_loss)
print('Validation: ')
print('[Epoch: %d, numImages: %5d]' % (epoch, i * self.args.batch_size ))
print("Acc:{}".format(acc), "mIoU:{}".format(mIoU))
print('Loss: %.3f' % val_loss)
# Save per epoch loss on loss.log
train_callbacks.append(
TrainLogger(args.log_interval,
train.nbatches(),
loss_log_path=join(log_path, f"loss.log")))
else:
print(f"[Logging: None]")
# Model saving configuration
print(f"[Model Saving: {args.save}]")
if args.save:
test_callbacks.append(BestModelSaver(log_path, args.loss))
# Evaluation configuration
metric = KNNF1ScoreMetric(config.test_distance, neighbors=10)
dev_evaluator = Evaluator(dev, metric, 'dev', test_callbacks)
test_evaluator = Evaluator(
test,
metric,
'test',
callbacks=[MetricFileLogger(log_path=join(log_path, 'test-metric.log'))])
train_callbacks.extend([dev_evaluator, test_evaluator])
# Training configuration
optim = config.optimizer(model, lr=args.lr)
model_loader = ModelLoader(
args.recover, args.recover_optim) if args.recover is not None else None
trainer = Trainer(args.loss,
model,
config.loss,
train,
def __init__(self, data_train, data_valid, image_base_dir, instructions):
"""
:param data_train:
:param data_valid:
:param image_base_dir:
:param instructions:
"""
self.image_base_dir = image_base_dir
self.data_valid = data_valid
self.instructions = instructions
# specify model save dir
self.model_name = instructions[STR.MODEL_NAME]
# now = time.localtime()
# start_time = "{}-{}-{}T{}:{}:{}".format(now.tm_year, now.tm_mon, now.tm_mday, now.tm_hour, now.tm_min,
# now.tm_sec)
experiment_folder_path = os.path.join(paths.MODELS_FOLDER_PATH,
self.model_name)
if os.path.exists(experiment_folder_path):
Warning(
"Experiment folder exists already. Files might be overwritten")
os.makedirs(experiment_folder_path, exist_ok=True)
# define saver and save instructions
self.saver = Saver(folder_path=experiment_folder_path,
instructions=instructions)
self.saver.save_instructions()
# define Tensorboard Summary
self.writer = SummaryWriter(log_dir=experiment_folder_path)
nn_input_size = instructions[STR.NN_INPUT_SIZE]
state_dict_file_path = instructions.get(STR.STATE_DICT_FILE_PATH, None)
self.colour_mapping = mapping.get_colour_mapping()
# define transformers for training
crops_per_image = instructions.get(STR.CROPS_PER_IMAGE, 10)
apply_random_cropping = (STR.CROPS_PER_IMAGE in instructions.keys()) and \
(STR.IMAGES_PER_BATCH in instructions.keys())
print("{}applying random cropping".format(
"" if apply_random_cropping else "_NOT_ "))
t = [Normalize()]
if apply_random_cropping:
t.append(
RandomCrop(min_size=instructions.get(STR.CROP_SIZE_MIN, 400),
max_size=instructions.get(STR.CROP_SIZE_MAX, 1000),
crop_count=crops_per_image))
t += [
Resize(nn_input_size),
Flip(p_vertical=0.2, p_horizontal=0.5),
ToTensor()
]
transformations_train = transforms.Compose(t)
# define transformers for validation
transformations_valid = transforms.Compose(
[Normalize(), Resize(nn_input_size),
ToTensor()])
# set up data loaders
dataset_train = DictArrayDataSet(image_base_dir=image_base_dir,
data=data_train,
num_classes=len(
self.colour_mapping.keys()),
transformation=transformations_train)
# define batch sizes
self.batch_size = instructions[STR.BATCH_SIZE]
if apply_random_cropping:
self.data_loader_train = DataLoader(
dataset=dataset_train,
batch_size=instructions[STR.IMAGES_PER_BATCH],
shuffle=True,
collate_fn=custom_collate)
else:
self.data_loader_train = DataLoader(dataset=dataset_train,
batch_size=self.batch_size,
shuffle=True,
collate_fn=custom_collate)
dataset_valid = DictArrayDataSet(image_base_dir=image_base_dir,
data=data_valid,
num_classes=len(
self.colour_mapping.keys()),
transformation=transformations_valid)
self.data_loader_valid = DataLoader(dataset=dataset_valid,
batch_size=self.batch_size,
shuffle=False,
collate_fn=custom_collate)
self.num_classes = dataset_train.num_classes()
# define model
print("Building model")
self.model = DeepLab(num_classes=self.num_classes,
backbone=instructions.get(STR.BACKBONE, "resnet"),
output_stride=instructions.get(
STR.DEEPLAB_OUTPUT_STRIDE, 16))
# load weights
if state_dict_file_path is not None:
print("loading state_dict from:")
print(state_dict_file_path)
load_state_dict(self.model, state_dict_file_path)
learning_rate = instructions.get(STR.LEARNING_RATE, 1e-5)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': learning_rate
}, {
'params': self.model.get_10x_lr_params(),
'lr': learning_rate
}]
# choose gpu or cpu
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
if instructions.get(STR.MULTI_GPU, False):
if torch.cuda.device_count() > 1:
print("Using ", torch.cuda.device_count(), " GPUs!")
self.model = nn.DataParallel(self.model)
self.model.to(self.device)
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=0.9,
weight_decay=5e-4,
nesterov=False)
# calculate class weights
if instructions.get(STR.CLASS_STATS_FILE_PATH, None):
class_weights = calculate_class_weights(
instructions[STR.CLASS_STATS_FILE_PATH],
self.colour_mapping,
modifier=instructions.get(STR.LOSS_WEIGHT_MODIFIER, 1.01))
class_weights = torch.from_numpy(class_weights.astype(np.float32))
else:
class_weights = None
self.criterion = SegmentationLosses(
weight=class_weights, cuda=self.device.type != "cpu").build_loss()
# Define Evaluator
self.evaluator = Evaluator(self.num_classes)
# Define lr scheduler
self.scheduler = None
if instructions.get(STR.USE_LR_SCHEDULER, True):
self.scheduler = LR_Scheduler(mode="cos",
base_lr=learning_rate,
num_epochs=instructions[STR.EPOCHS],
iters_per_epoch=len(
self.data_loader_train))
# print information before training start
print("-" * 60)
print("instructions")
pprint(instructions)
model_parameters = sum([p.nelement() for p in self.model.parameters()])
print("Model parameters: {:.2E}".format(model_parameters))
self.best_prediction = 0.0
class Trainer:
def __init__(self, data_train, data_valid, image_base_dir, instructions):
"""
:param data_train:
:param data_valid:
:param image_base_dir:
:param instructions:
"""
self.image_base_dir = image_base_dir
self.data_valid = data_valid
self.instructions = instructions
# specify model save dir
self.model_name = instructions[STR.MODEL_NAME]
# now = time.localtime()
# start_time = "{}-{}-{}T{}:{}:{}".format(now.tm_year, now.tm_mon, now.tm_mday, now.tm_hour, now.tm_min,
# now.tm_sec)
experiment_folder_path = os.path.join(paths.MODELS_FOLDER_PATH,
self.model_name)
if os.path.exists(experiment_folder_path):
Warning(
"Experiment folder exists already. Files might be overwritten")
os.makedirs(experiment_folder_path, exist_ok=True)
# define saver and save instructions
self.saver = Saver(folder_path=experiment_folder_path,
instructions=instructions)
self.saver.save_instructions()
# define Tensorboard Summary
self.writer = SummaryWriter(log_dir=experiment_folder_path)
nn_input_size = instructions[STR.NN_INPUT_SIZE]
state_dict_file_path = instructions.get(STR.STATE_DICT_FILE_PATH, None)
self.colour_mapping = mapping.get_colour_mapping()
# define transformers for training
crops_per_image = instructions.get(STR.CROPS_PER_IMAGE, 10)
apply_random_cropping = (STR.CROPS_PER_IMAGE in instructions.keys()) and \
(STR.IMAGES_PER_BATCH in instructions.keys())
print("{}applying random cropping".format(
"" if apply_random_cropping else "_NOT_ "))
t = [Normalize()]
if apply_random_cropping:
t.append(
RandomCrop(min_size=instructions.get(STR.CROP_SIZE_MIN, 400),
max_size=instructions.get(STR.CROP_SIZE_MAX, 1000),
crop_count=crops_per_image))
t += [
Resize(nn_input_size),
Flip(p_vertical=0.2, p_horizontal=0.5),
ToTensor()
]
transformations_train = transforms.Compose(t)
# define transformers for validation
transformations_valid = transforms.Compose(
[Normalize(), Resize(nn_input_size),
ToTensor()])
# set up data loaders
dataset_train = DictArrayDataSet(image_base_dir=image_base_dir,
data=data_train,
num_classes=len(
self.colour_mapping.keys()),
transformation=transformations_train)
# define batch sizes
self.batch_size = instructions[STR.BATCH_SIZE]
if apply_random_cropping:
self.data_loader_train = DataLoader(
dataset=dataset_train,
batch_size=instructions[STR.IMAGES_PER_BATCH],
shuffle=True,
collate_fn=custom_collate)
else:
self.data_loader_train = DataLoader(dataset=dataset_train,
batch_size=self.batch_size,
shuffle=True,
collate_fn=custom_collate)
dataset_valid = DictArrayDataSet(image_base_dir=image_base_dir,
data=data_valid,
num_classes=len(
self.colour_mapping.keys()),
transformation=transformations_valid)
self.data_loader_valid = DataLoader(dataset=dataset_valid,
batch_size=self.batch_size,
shuffle=False,
collate_fn=custom_collate)
self.num_classes = dataset_train.num_classes()
# define model
print("Building model")
self.model = DeepLab(num_classes=self.num_classes,
backbone=instructions.get(STR.BACKBONE, "resnet"),
output_stride=instructions.get(
STR.DEEPLAB_OUTPUT_STRIDE, 16))
# load weights
if state_dict_file_path is not None:
print("loading state_dict from:")
print(state_dict_file_path)
load_state_dict(self.model, state_dict_file_path)
learning_rate = instructions.get(STR.LEARNING_RATE, 1e-5)
train_params = [{
'params': self.model.get_1x_lr_params(),
'lr': learning_rate
}, {
'params': self.model.get_10x_lr_params(),
'lr': learning_rate
}]
# choose gpu or cpu
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
if instructions.get(STR.MULTI_GPU, False):
if torch.cuda.device_count() > 1:
print("Using ", torch.cuda.device_count(), " GPUs!")
self.model = nn.DataParallel(self.model)
self.model.to(self.device)
# Define Optimizer
self.optimizer = torch.optim.SGD(train_params,
momentum=0.9,
weight_decay=5e-4,
nesterov=False)
# calculate class weights
if instructions.get(STR.CLASS_STATS_FILE_PATH, None):
class_weights = calculate_class_weights(
instructions[STR.CLASS_STATS_FILE_PATH],
self.colour_mapping,
modifier=instructions.get(STR.LOSS_WEIGHT_MODIFIER, 1.01))
class_weights = torch.from_numpy(class_weights.astype(np.float32))
else:
class_weights = None
self.criterion = SegmentationLosses(
weight=class_weights, cuda=self.device.type != "cpu").build_loss()
# Define Evaluator
self.evaluator = Evaluator(self.num_classes)
# Define lr scheduler
self.scheduler = None
if instructions.get(STR.USE_LR_SCHEDULER, True):
self.scheduler = LR_Scheduler(mode="cos",
base_lr=learning_rate,
num_epochs=instructions[STR.EPOCHS],
iters_per_epoch=len(
self.data_loader_train))
# print information before training start
print("-" * 60)
print("instructions")
pprint(instructions)
model_parameters = sum([p.nelement() for p in self.model.parameters()])
print("Model parameters: {:.2E}".format(model_parameters))
self.best_prediction = 0.0
def train(self, epoch):
self.model.train()
train_loss = 0.0
# create a progress bar
pbar = tqdm(self.data_loader_train)
num_batches_train = len(self.data_loader_train)
# go through each item in the training data
for i, sample in enumerate(pbar):
# set input and target
nn_input = sample[STR.NN_INPUT].to(self.device)
nn_target = sample[STR.NN_TARGET].to(self.device,
dtype=torch.float)
if self.scheduler:
self.scheduler(self.optimizer, i, epoch, self.best_prediction)
# run model
output = self.model(nn_input)
# calc losses
loss = self.criterion(output, nn_target)
# # save step losses
# combined_loss_steps.append(float(loss))
# regression_loss_steps.append(float(regression_loss))
# classification_loss_steps.append(float(classification_loss))
train_loss += loss.item()
pbar.set_description('Train loss: %.3f' % (train_loss / (i + 1)))
self.writer.add_scalar('train/total_loss_iter', loss.item(),
i + num_batches_train * epoch)
# calculate gradient and update model weights
loss.backward()
# torch.nn.utils.clip_grad_norm_(model.parameters(), 0.1)
self.optimizer.step()
self.optimizer.zero_grad()
self.writer.add_scalar('train/total_loss_epoch', train_loss, epoch)
print("[Epoch: {}, num images/crops: {}]".format(
epoch, num_batches_train * self.batch_size))
print("Loss: {:.2f}".format(train_loss))
def validation(self, epoch):
self.model.eval()
self.evaluator.reset()
test_loss = 0.0
pbar = tqdm(self.data_loader_valid, desc='\r')
num_batches_val = len(self.data_loader_valid)
for i, sample in enumerate(pbar):
# set input and target
nn_input = sample[STR.NN_INPUT].to(self.device)
nn_target = sample[STR.NN_TARGET].to(self.device,
dtype=torch.float)
with torch.no_grad():
output = self.model(nn_input)
loss = self.criterion(output, nn_target)
test_loss += loss.item()
pbar.set_description('Test loss: %.3f' % (test_loss / (i + 1)))
pred = output.data.cpu().numpy()
pred = np.argmax(pred, axis=1)
nn_target = nn_target.cpu().numpy()
# Add batch sample into evaluator
self.evaluator.add_batch(nn_target, pred)
# Fast test during the training
Acc = self.evaluator.Pixel_Accuracy()
Acc_class = self.evaluator.Pixel_Accuracy_Class()
mIoU = self.evaluator.Mean_Intersection_over_Union()
FWIoU = self.evaluator.Frequency_Weighted_Intersection_over_Union()
self.writer.add_scalar('val/total_loss_epoch', test_loss, epoch)
self.writer.add_scalar('val/mIoU', mIoU, epoch)
self.writer.add_scalar('val/Acc', Acc, epoch)
self.writer.add_scalar('val/Acc_class', Acc_class, epoch)
self.writer.add_scalar('val/fwIoU', FWIoU, epoch)
print('Validation:')
print("[Epoch: {}, num crops: {}]".format(
epoch, num_batches_val * self.batch_size))
print(
"Acc:{:.2f}, Acc_class:{:.2f}, mIoU:{:.2f}, fwIoU: {:.2f}".format(
Acc, Acc_class, mIoU, FWIoU))
print("Loss: {:.2f}".format(test_loss))
new_pred = mIoU
is_best = new_pred > self.best_prediction
if is_best:
self.best_prediction = new_pred
self.saver.save_checkpoint(self.model, is_best, epoch)
class PPO(object):
def __init__(self,
env_cls,
args,
manager,
cfg,
process_num,
pre=False,
pre_irl=False,
infer=False):
"""
:param env_cls: env class or function, not instance, as we need to create several instance in class.
:param args:
:param manager:
:param cfg:
:param process_num: process number
:param pre: set to pretrain mode
:param infer: set to test mode
"""
self.process_num = process_num
# initialize envs for each process
self.env_list = []
for _ in range(process_num):
self.env_list.append(env_cls())
# construct policy and value network
self.policy = MultiDiscretePolicy(cfg).to(device=DEVICE)
self.value = Value(cfg).to(device=DEVICE)
if pre:
self.print_per_batch = args.print_per_batch
from dbquery import DBQuery
db = DBQuery(args.data_dir)
self.data_train = manager.create_dataset_rl(
'train', args.batchsz, cfg, db)
self.data_valid = manager.create_dataset_rl(
'valid', args.batchsz, cfg, db)
self.data_test = manager.create_dataset_rl('test', args.batchsz,
cfg, db)
self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
else:
self.rewarder = RewardEstimator(args,
manager,
cfg,
pretrain=pre_irl,
inference=infer)
self.evaluator = Evaluator(args.data_dir, cfg)
self.save_dir = args.save_dir
self.save_per_epoch = args.save_per_epoch
self.optim_batchsz = args.batchsz
self.update_round = args.update_round
self.policy.eval()
self.value.eval()
self.gamma = args.gamma
self.epsilon = args.epsilon
self.tau = args.tau
self.policy_optim = optim.RMSprop(self.policy.parameters(),
lr=args.lr_rl)
self.value_optim = optim.Adam(self.value.parameters(), lr=args.lr_rl)
def policy_loop(self, data):
s, target_a = to_device(data)
a_weights = self.policy(s)
loss_a = self.multi_entropy_loss(a_weights, target_a)
return loss_a
def imitating(self, epoch):
"""
pretrain the policy by simple imitation learning (behavioral cloning)
"""
self.policy.train()
a_loss = 0.
for i, data in enumerate(self.data_train):
self.policy_optim.zero_grad()
loss_a = self.policy_loop(data)
a_loss += loss_a.item()
loss_a.backward()
self.policy_optim.step()
if (i + 1) % self.print_per_batch == 0:
a_loss /= self.print_per_batch
logging.debug(
'<<dialog policy>> epoch {}, iter {}, loss_a:{}'.format(
epoch, i, a_loss))
a_loss = 0.
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch, True)
self.policy.eval()
def imit_test(self, epoch, best):
"""
provide an unbiased evaluation of the policy fit on the training dataset
"""
a_loss = 0.
for i, data in enumerate(self.data_valid):
loss_a = self.policy_loop(data)
a_loss += loss_a.item()
a_loss /= len(self.data_valid)
logging.debug(
'<<dialog policy>> validation, epoch {}, loss_a:{}'.format(
epoch, a_loss))
if a_loss < best:
logging.info('<<dialog policy>> best model saved')
best = a_loss
self.save(self.save_dir, 'best', True)
a_loss = 0.
for i, data in enumerate(self.data_test):
loss_a = self.policy_loop(data)
a_loss += loss_a.item()
a_loss /= len(self.data_test)
logging.debug('<<dialog policy>> test, epoch {}, loss_a:{}'.format(
epoch, a_loss))
return best
def imit_value(self, epoch, batchsz, best):
self.value.train()
batch = self.sample(batchsz)
s = torch.from_numpy(np.stack(batch.state)).to(device=DEVICE)
a = torch.from_numpy(np.stack(batch.action)).to(device=DEVICE)
next_s = torch.from_numpy(np.stack(batch.next_state)).to(device=DEVICE)
mask = torch.Tensor(np.stack(batch.mask)).to(device=DEVICE)
batchsz = s.size(0)
v = self.value(s).squeeze(-1).detach()
log_pi_old_sa = self.policy.get_log_prob(s, a).detach()
r = self.rewarder.estimate(s, a, next_s, log_pi_old_sa).detach()
A_sa, v_target = self.est_adv(r, v, mask)
for i in range(self.update_round):
perm = torch.randperm(batchsz)
v_target_shuf, s_shuf = v_target[perm], s[perm]
optim_chunk_num = int(np.ceil(batchsz / self.optim_batchsz))
v_target_shuf, s_shuf = torch.chunk(v_target_shuf,
optim_chunk_num), torch.chunk(
s_shuf, optim_chunk_num)
value_loss = 0.
for v_target_b, s_b in zip(v_target_shuf, s_shuf):
self.value_optim.zero_grad()
v_b = self.value(s_b).squeeze(-1)
loss = (v_b - v_target_b).pow(2).mean()
value_loss += loss.item()
loss.backward()
self.value_optim.step()
value_loss /= optim_chunk_num
logging.debug(
'<<dialog policy>> epoch {}, iteration {}, loss {}'.format(
epoch, i, value_loss))
if value_loss < best:
logging.info('<<dialog policy>> best model saved')
best = value_loss
self.save(self.save_dir, 'best', True)
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch, True)
self.value.eval()
return best
def train_irl(self, epoch, batchsz):
batch = self.sample(batchsz)
self.rewarder.train_irl(batch, epoch)
def test_irl(self, epoch, batchsz, best):
batch = self.sample(batchsz)
best = self.rewarder.test_irl(batch, epoch, best)
return best
def est_adv(self, r, v, mask):
"""
we save a trajectory in continuous space and it reaches the ending of current trajectory when mask=0.
:param r: reward, Tensor, [b]
:param v: estimated value, Tensor, [b]
:param mask: indicates ending for 0 otherwise 1, Tensor, [b]
:return: A(s, a), V-target(s), both Tensor
"""
batchsz = v.size(0)
# v_target is worked out by Bellman equation.
v_target = torch.Tensor(batchsz).to(device=DEVICE)
delta = torch.Tensor(batchsz).to(device=DEVICE)
A_sa = torch.Tensor(batchsz).to(device=DEVICE)
prev_v_target = 0
prev_v = 0
prev_A_sa = 0
for t in reversed(range(batchsz)):
# mask here indicates a end of trajectory
# this value will be treated as the target value of value network.
# mask = 0 means the immediate reward is the real V(s) since it's end of trajectory.
# formula: V(s_t) = r_t + gamma * V(s_t+1)
v_target[t] = r[t] + self.gamma * prev_v_target * mask[t]
# please refer to : https://arxiv.org/abs/1506.02438
# for generalized adavantage estimation
# formula: delta(s_t) = r_t + gamma * V(s_t+1) - V(s_t)
delta[t] = r[t] + self.gamma * prev_v * mask[t] - v[t]
# formula: A(s, a) = delta(s_t) + gamma * lamda * A(s_t+1, a_t+1)
# here use symbol tau as lambda, but original paper uses symbol lambda.
A_sa[t] = delta[t] + self.gamma * self.tau * prev_A_sa * mask[t]
# update previous
prev_v_target = v_target[t]
prev_v = v[t]
prev_A_sa = A_sa[t]
# normalize A_sa
A_sa = (A_sa - A_sa.mean()) / A_sa.std()
return A_sa, v_target
def update(self, batchsz, epoch, best=None):
"""
firstly sample batchsz items and then perform optimize algorithms.
:param batchsz:
:param epoch:
:param best:
:return:
"""
backward = True if best is None else False
if backward:
self.policy.train()
self.value.train()
# 1. sample data asynchronously
batch = self.sample(batchsz)
# data in batch is : batch.state: ([1, s_dim], [1, s_dim]...)
# batch.action: ([1, a_dim], [1, a_dim]...)
# batch.reward/ batch.mask: ([1], [1]...)
s = torch.from_numpy(np.stack(batch.state)).to(device=DEVICE)
a = torch.from_numpy(np.stack(batch.action)).to(device=DEVICE)
next_s = torch.from_numpy(np.stack(batch.next_state)).to(device=DEVICE)
mask = torch.Tensor(np.stack(batch.mask)).to(device=DEVICE)
batchsz = s.size(0)
# 2. update reward estimator
inputs = (s, a, next_s)
if backward:
self.rewarder.update_irl(inputs, batchsz, epoch)
else:
best[1] = self.rewarder.update_irl(inputs, batchsz, epoch, best[1])
# 3. get estimated V(s) and PI_old(s, a)
# actually, PI_old(s, a) can be saved when interacting with env, so as to save the time of one forward elapsed
# v: [b, 1] => [b]
v = self.value(s).squeeze(-1).detach()
log_pi_old_sa = self.policy.get_log_prob(s, a).detach()
# 4. estimate advantage and v_target according to GAE and Bellman Equation
r = self.rewarder.estimate(s, a, next_s, log_pi_old_sa).detach()
A_sa, v_target = self.est_adv(r, v, mask)
if backward:
logging.debug('<<dialog policy>> epoch {}, reward {}'.format(
epoch,
r.mean().item()))
else:
reward = r.mean().item()
logging.debug(
'<<dialog policy>> validation, epoch {}, reward {}'.format(
epoch, reward))
if reward > best[2]:
logging.info('<<dialog policy>> best model saved')
best[2] = reward
self.save(self.save_dir, 'best', True)
with open(self.save_dir + '/best.pkl', 'wb') as f:
pickle.dump(best, f)
return best
# 5. update dialog policy
for i in range(self.update_round):
# 1. shuffle current batch
perm = torch.randperm(batchsz)
# shuffle the variable for mutliple optimize
v_target_shuf, A_sa_shuf, s_shuf, a_shuf, log_pi_old_sa_shuf = v_target[perm], A_sa[perm], s[perm], a[perm], \
log_pi_old_sa[perm]
# 2. get mini-batch for optimizing
optim_chunk_num = int(np.ceil(batchsz / self.optim_batchsz))
# chunk the optim_batch for total batch
v_target_shuf, A_sa_shuf, s_shuf, a_shuf, log_pi_old_sa_shuf = torch.chunk(v_target_shuf, optim_chunk_num), \
torch.chunk(A_sa_shuf, optim_chunk_num), \
torch.chunk(s_shuf, optim_chunk_num), \
torch.chunk(a_shuf, optim_chunk_num), \
torch.chunk(log_pi_old_sa_shuf,
optim_chunk_num)
# 3. iterate all mini-batch to optimize
policy_loss, value_loss = 0., 0.
for v_target_b, A_sa_b, s_b, a_b, log_pi_old_sa_b in zip(
v_target_shuf, A_sa_shuf, s_shuf, a_shuf,
log_pi_old_sa_shuf):
# print('optim:', batchsz, v_target_b.size(), A_sa_b.size(), s_b.size(), a_b.size(), log_pi_old_sa_b.size())
# 1. update value network
self.value_optim.zero_grad()
v_b = self.value(s_b).squeeze(-1)
loss = (v_b - v_target_b).pow(2).mean()
value_loss += loss.item()
# backprop
loss.backward()
# nn.utils.clip_grad_norm(self.value.parameters(), 4)
self.value_optim.step()
# 2. update policy network by clipping
self.policy_optim.zero_grad()
# [b, 1]
log_pi_sa = self.policy.get_log_prob(s_b, a_b)
# ratio = exp(log_Pi(a|s) - log_Pi_old(a|s)) = Pi(a|s) / Pi_old(a|s)
# we use log_pi for stability of numerical operation
# [b, 1] => [b]
ratio = (log_pi_sa - log_pi_old_sa_b).exp().squeeze(-1)
surrogate1 = ratio * A_sa_b
surrogate2 = torch.clamp(ratio, 1 - self.epsilon,
1 + self.epsilon) * A_sa_b
# this is element-wise comparing.
# we add negative symbol to convert gradient ascent to gradient descent
surrogate = -torch.min(surrogate1, surrogate2).mean()
policy_loss += surrogate.item()
# backprop
surrogate.backward()
# gradient clipping, for stability
torch.nn.utils.clip_grad_norm(self.policy.parameters(), 10)
# self.lock.acquire() # retain lock to update weights
self.policy_optim.step()
# self.lock.release() # release lock
value_loss /= optim_chunk_num
policy_loss /= optim_chunk_num
logging.debug(
'<<dialog policy>> epoch {}, iteration {}, value, loss {}'.
format(epoch, i, value_loss))
logging.debug(
'<<dialog policy>> epoch {}, iteration {}, policy, loss {}'.
format(epoch, i, policy_loss))
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch)
with open(self.save_dir + '/' + str(epoch) + '.pkl', 'wb') as f:
pickle.dump(best, f)
self.policy.eval()
self.value.eval()
def sample(self, batchsz):
"""
Given batchsz number of task, the batchsz will be splited equally to each processes
and when processes return, it merge all data and return
:param batchsz:
:return: batch
"""
# batchsz will be splitted into each process,
# final batchsz maybe larger than batchsz parameters
process_batchsz = np.ceil(batchsz / self.process_num).astype(np.int32)
# buffer to save all data
queue = mp.Queue()
# start processes for pid in range(1, processnum)
# if processnum = 1, this part will be ignored.
# when save tensor in Queue, the process should keep alive till Queue.get(),
# please refer to : https://discuss.pytorch.org/t/using-torch-tensor-over-multiprocessing-queue-process-fails/2847
# however still some problem on CUDA tensors on multiprocessing queue,
# please refer to : https://discuss.pytorch.org/t/cuda-tensors-on-multiprocessing-queue/28626
# so just transform tensors into numpy, then put them into queue.
evt = mp.Event()
processes = []
for i in range(self.process_num):
process_args = (i, queue, evt, self.env_list[i], self.policy,
process_batchsz)
processes.append(mp.Process(target=sampler, args=process_args))
for p in processes:
# set the process as daemon, and it will be killed once the main process is stoped.
p.daemon = True
p.start()
# we need to get the first Memory object and then merge others Memory use its append function.
pid0, buff0 = queue.get()
for _ in range(1, self.process_num):
pid, buff_ = queue.get()
buff0.append(buff_) # merge current Memory into buff0
evt.set()
# now buff saves all the sampled data
buff = buff0
return buff.get_batch()
def evaluate(self):
env = self.env_list[0]
traj_len = 40
reward_tot, turn_tot, inform_tot, match_tot, success_tot = [], [], [], [], []
for seed in range(1000):
s = env.reset(seed)
print('seed', seed)
print('goal', env.goal.domain_goals)
print('usr', s['user_action'])
turn = traj_len
reward = []
value = []
mask = []
for t in range(traj_len):
s_vec = torch.Tensor(state_vectorize(s, env.cfg,
env.db)).to(device=DEVICE)
# mode with policy during evaluation
a = self.policy.select_action(s_vec, False)
next_s, done = env.step(s, a.cpu())
next_s_vec = torch.Tensor(
state_vectorize(next_s, env.cfg, env.db)).to(device=DEVICE)
log_pi = self.policy.get_log_prob(s_vec, a)
r = self.rewarder.estimate(s_vec, a, next_s_vec, log_pi)
v = self.value(s_vec).squeeze(-1)
reward.append(r.item())
value.append(v.item())
s = next_s
print('sys', s['last_sys_action'])
print('usr', s['user_action'])
if done:
mask.append(0)
turn = t + 2 # one due to counting from 0, the one for the last turn
break
mask.append(1)
reward_tot.append(np.mean(reward))
turn_tot.append(turn)
match_tot += self.evaluator.match_rate(s)
inform_tot.append(self.evaluator.inform_F1(s))
reward = torch.Tensor(reward)
value = torch.Tensor(value)
mask = torch.LongTensor(mask)
A_sa, v_target = self.est_adv(reward, value, mask)
print('turn', turn)
#print('reward', A_sa.tolist())
print('reward', v_target[0].item())
match_session = self.evaluator.match_rate(s, True)
print('match', match_session)
inform_session = self.evaluator.inform_F1(s, True)
print('inform', inform_session)
if (match_session == 1 and inform_session[1] == 1) \
or (match_session == 1 and inform_session[1] is None) \
or (match_session is None and inform_session[1] == 1):
print('success', 1)
success_tot.append(1)
else:
print('success', 0)
success_tot.append(0)
logging.info('reward {}'.format(np.mean(reward_tot)))
logging.info('turn {}'.format(np.mean(turn_tot)))
logging.info('match {}'.format(np.mean(match_tot)))
TP, FP, FN = np.sum(inform_tot, 0)
prec = TP / (TP + FP)
rec = TP / (TP + FN)
F1 = 2 * prec * rec / (prec + rec)
logging.info('inform rec {}, F1 {}'.format(rec, F1))
logging.info('success {}'.format(np.mean(success_tot)))
def save(self, directory, epoch, rl_only=False):
if not os.path.exists(directory):
os.makedirs(directory)
if not rl_only:
self.rewarder.save_irl(directory, epoch)
torch.save(self.value.state_dict(),
directory + '/' + str(epoch) + '_ppo.val.mdl')
torch.save(self.policy.state_dict(),
directory + '/' + str(epoch) + '_ppo.pol.mdl')
logging.info(
'<<dialog policy>> epoch {}: saved network to mdl'.format(epoch))
def load(self, filename):
self.rewarder.load_irl(filename)
value_mdl = filename + '_ppo.val.mdl'
policy_mdl = filename + '_ppo.pol.mdl'
if os.path.exists(value_mdl):
self.value.load_state_dict(torch.load(value_mdl))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
value_mdl))
if os.path.exists(policy_mdl):
self.policy.load_state_dict(torch.load(policy_mdl))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
policy_mdl))
best_pkl = filename + '.pkl'
if os.path.exists(best_pkl):
with open(best_pkl, 'rb') as f:
best = pickle.load(f)
else:
best = [float('inf'), float('inf'), float('-inf')]
return best
def model_fn(features, labels, mode, params):
"""
This is a function for creating a computational tensorflow graph.
The function is in format required by tf.estimator.
"""
is_training = mode == tf.estimator.ModeKeys.TRAIN
# the base network
def backbone(images, is_training):
if params['backbone'] == 'mobilenet':
return mobilenet_v1(images,
is_training,
depth_multiplier=params['depth_multiplier'])
elif params['backbone'] == 'shufflenet':
return shufflenet_v2(images,
is_training,
depth_multiplier=str(
params['depth_multiplier']))
# add additional layers to the base network
feature_extractor = RetinaNetFeatureExtractor(is_training, backbone)
# ssd anchor maker
anchor_generator = AnchorGenerator(strides=[8, 16, 32, 64, 128],
scales=[32, 64, 128, 256, 512],
scale_multipliers=[1.0, 1.4142],
aspect_ratios=[1.0, 2.0, 0.5])
num_anchors_per_location = anchor_generator.num_anchors_per_location
# add layers that predict boxes and labels
box_predictor = RetinaNetBoxPredictor(is_training, params['num_classes'],
num_anchors_per_location)
# collect everything on one place
ssd = SSD(features['images'], feature_extractor, anchor_generator,
box_predictor, params['num_classes'])
# add nms to the graph
if not is_training:
predictions = ssd.get_predictions(
score_threshold=params['score_threshold'],
iou_threshold=params['iou_threshold'],
max_boxes_per_class=params['max_boxes_per_class'])
if mode == tf.estimator.ModeKeys.PREDICT:
# because images are resized before
# feeding them to the network
box_scaler = features['box_scaler']
predictions['boxes'] /= box_scaler
export_outputs = tf.estimator.export.PredictOutput({
name: tf.identity(tensor, name)
for name, tensor in predictions.items()
})
return tf.estimator.EstimatorSpec(
mode,
predictions=predictions,
export_outputs={'outputs': export_outputs})
# add l2 regularization
with tf.name_scope('weight_decay'):
add_weight_decay(params['weight_decay'])
regularization_loss = tf.losses.get_regularization_loss()
# create localization and classification losses
losses = ssd.loss(labels, params)
tf.losses.add_loss(params['localization_loss_weight'] *
losses['localization_loss'])
tf.losses.add_loss(params['classification_loss_weight'] *
losses['classification_loss'])
tf.summary.scalar('regularization_loss', regularization_loss)
tf.summary.scalar('localization_loss', losses['localization_loss'])
tf.summary.scalar('classification_loss', losses['classification_loss'])
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
if mode == tf.estimator.ModeKeys.EVAL:
batch_size = features['images'].shape[0].value
assert batch_size == 1
evaluator = Evaluator(num_classes=params['num_classes'])
eval_metric_ops = evaluator.get_metric_ops(labels, predictions)
return tf.estimator.EstimatorSpec(mode,
loss=total_loss,
eval_metric_ops=eval_metric_ops)
assert mode == tf.estimator.ModeKeys.TRAIN
with tf.variable_scope('learning_rate'):
global_step = tf.train.get_global_step()
learning_rate = tf.train.cosine_decay(params['initial_learning_rate'],
global_step,
decay_steps=params['num_steps'])
tf.summary.scalar('learning_rate', learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops), tf.variable_scope('optimizer'):
optimizer = tf.train.AdamOptimizer(learning_rate)
grads_and_vars = optimizer.compute_gradients(total_loss)
train_op = optimizer.apply_gradients(grads_and_vars, global_step)
for g, v in grads_and_vars:
tf.summary.histogram(v.name[:-2] + '_hist', v)
tf.summary.histogram(v.name[:-2] + '_grad_hist', g)
with tf.control_dependencies([train_op]), tf.name_scope('ema'):
ema = tf.train.ExponentialMovingAverage(decay=MOVING_AVERAGE_DECAY,
num_updates=global_step)
train_op = ema.apply(tf.trainable_variables())
return tf.estimator.EstimatorSpec(mode, loss=total_loss, train_op=train_op)
weight = np.load(classes_weights_path)
else:
weight = calculate_weigths_labels(dataset, dataloader, 18)
weight = torch.from_numpy(weight.astype(np.float32))
"""
5.2、定义损失函数
"""
criterion = SegmentationLosses(weight=weight, cuda=False).build_loss(mode="ce")
"""
6、定义评价指标
"""
evaluator = Evaluator(18)
"""
7、定义学习率调度器
"""
scheduler = LR_Scheduler(lr_scheduler, lr, epochs, len(dataloader))
"""
8、使用预定义模型，并判断是否使用cuda
if os.path.isdir(model_save):
try:
checkpoint = torch.load(model_save+'model.t7', map_location='cpu')
model.load_state_dict(checkpoint['state'])
start_epoch = checkpoint['epoch']
print('===> Load last checkpoint data')
def main():
nclass = 7
lr = 0.01
num_epochs = 400
batch_size = 4
best_pred = 0.0
drop_last = True
train_loader, val_loader = make_dataloader(batch_size=batch_size,
drop_last=drop_last)
model = DeepLab(num_classes=nclass, output_stride=16, freeze_bn=False)
train_params = [{
'params': model.get_1x_lr_params(),
'lr': lr
}, {
'params': model.get_10x_lr_params(),
'lr': lr * 10
}]
optimizer = torch.optim.SGD(train_params,
momentum=0.9,
weight_decay=5e-4,
nesterov=False)
criterion = SegmentationLosses(weight=None,
cuda=True).build_loss(mode='ce')
evaluator = Evaluator(nclass)
scheduler = LR_Scheduler(mode='poly',
base_lr=lr,
num_epochs=num_epochs,
iters_per_epoch=len(train_loader))
model = model.cuda()
for epoch in range(num_epochs):
train_loss = 0.0
count = 0
model.train()
for i, sample in enumerate(train_loader):
image, target = sample['image'], sample['label']
image, target = image.cuda(), target.cuda()
scheduler(optimizer, i, epoch, best_pred)
optimizer.zero_grad()
output = model(image)
loss = criterion(output, target)
loss.backward()
optimizer.step()
train_loss += loss.item()
count += 1
train_loss = train_loss / count
print(
'Training [Epoch: %d, best_pred: %.4f, numImages: %5d, Loss: %.3f]'
% (epoch, best_pred, i * batch_size + image.data.shape[0],
train_loss))
writer.add_scalar('scalar/loss_seg_train', train_loss, epoch)
model.eval()
evaluator.reset()
test_loss = 0.0
count = 0
for i, sample in enumerate(val_loader):
image, target = sample['image'], sample['label']
image, target = image.cuda(), target.cuda()
with torch.no_grad():
output = model(image)
loss = criterion(output, target)
test_loss += loss.item()
count += 1
pred = output.data.cpu().numpy()
target = target.cpu().numpy()
pred = np.argmax(pred, axis=1)
# Add batch sample into evaluator
evaluator.add_batch(target, pred)
# Fast test during the training
test_loss = test_loss / count
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
print('Validation [Epoch: %d, numImages: %5d, Loss: %.3f]' %
(epoch, i * batch_size + image.data.shape[0], test_loss))
print("Acc:{}, Acc_class:{}, mIoU:{}, fwIoU: {}".format(
Acc, Acc_class, mIoU, FWIoU))
writer.add_scalar('scalar/loss_seg_val', test_loss, epoch)
writer.add_scalar('scalar/Acc_val', Acc, epoch)
writer.add_scalar('scalar/Acc_class_val', Acc_class, epoch)
writer.add_scalar('scalar/mIou_val', mIoU, epoch)
writer.add_scalar('scalar/FWIou_val', FWIoU, epoch)
path = '/home/user/'
if mIoU > best_pred:
best_pred = mIoU
torch.save(model.state_dict(), path + 'model_params.pkl')
print('save the segmentation ' + str(epoch) +
'model, replace the previous parameters')
class Trainer:
def __init__(self, settings):
self.settings = settings
self.device = torch.device(
"cuda" if len(self.settings.gpu_devices) else "cpu")
self.gpu_devices = ','.join([str(id) for id in settings.gpu_devices])
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu_devices
self.log_path = os.path.join(self.settings.log_dir,
self.settings.model_name)
# checking the input height and width are multiples of 32
assert self.settings.height % 32 == 0, "input height must be a multiple of 32"
assert self.settings.width % 32 == 0, "input width must be a multiple of 32"
# data
datasets_dict = {
"3d60": datasets.ThreeD60,
"panosuncg": datasets.PanoSunCG,
"stanford2d3d": datasets.Stanford2D3D,
"matterport3d": datasets.Matterport3D
}
self.dataset = datasets_dict[self.settings.dataset]
self.settings.cube_w = self.settings.height // 2
fpath = os.path.join(os.path.dirname(__file__), "datasets",
"{}_{}.txt")
train_file_list = fpath.format(self.settings.dataset, "train")
val_file_list = fpath.format(self.settings.dataset, "val")
train_dataset = self.dataset(
self.settings.data_path,
train_file_list,
self.settings.height,
self.settings.width,
self.settings.disable_color_augmentation,
self.settings.disable_LR_filp_augmentation,
self.settings.disable_yaw_rotation_augmentation,
is_training=True)
self.train_loader = DataLoader(train_dataset,
self.settings.batch_size,
True,
num_workers=self.settings.num_workers,
pin_memory=True,
drop_last=True)
num_train_samples = len(train_dataset)
self.num_total_steps = num_train_samples // self.settings.batch_size * self.settings.num_epochs
val_dataset = self.dataset(
self.settings.data_path,
val_file_list,
self.settings.height,
self.settings.width,
self.settings.disable_color_augmentation,
self.settings.disable_LR_filp_augmentation,
self.settings.disable_yaw_rotation_augmentation,
is_training=False)
self.val_loader = DataLoader(val_dataset,
self.settings.batch_size,
False,
num_workers=self.settings.num_workers,
pin_memory=True,
drop_last=True)
# network
Net_dict = {"UniFuse": UniFuse, "Equi": Equi}
Net = Net_dict[self.settings.net]
self.model = Net(self.settings.num_layers,
self.settings.height,
self.settings.width,
self.settings.imagenet_pretrained,
train_dataset.max_depth_meters,
fusion_type=self.settings.fusion,
se_in_fusion=self.settings.se_in_fusion)
self.model.to(self.device)
self.parameters_to_train = list(self.model.parameters())
self.optimizer = optim.Adam(self.parameters_to_train,
self.settings.learning_rate)
if self.settings.load_weights_dir is not None:
self.load_model()
print("Training model named:\n ", self.settings.model_name)
print("Models and tensorboard events files are saved to:\n",
self.settings.log_dir)
print("Training is using:\n ", self.device)
self.compute_loss = BerhuLoss()
self.evaluator = Evaluator()
self.writers = {}
for mode in ["train", "val"]:
self.writers[mode] = SummaryWriter(
os.path.join(self.log_path, mode))
self.save_settings()
def train(self):
"""Run the entire training pipeline
"""
self.epoch = 0
self.step = 0
self.start_time = time.time()
self.validate()
for self.epoch in range(self.settings.num_epochs):
self.train_one_epoch()
self.validate()
if (self.epoch + 1) % self.settings.save_frequency == 0:
self.save_model()
def train_one_epoch(self):
"""Run a single epoch of training
"""
self.model.train()
pbar = tqdm.tqdm(self.train_loader)
pbar.set_description("Training Epoch_{}".format(self.epoch))
for batch_idx, inputs in enumerate(pbar):
outputs, losses = self.process_batch(inputs)
self.optimizer.zero_grad()
losses["loss"].backward()
self.optimizer.step()
# log less frequently after the first 1000 steps to save time & disk space
early_phase = batch_idx % self.settings.log_frequency == 0 and self.step < 1000
late_phase = self.step % 1000 == 0
if early_phase or late_phase:
pred_depth = outputs["pred_depth"].detach()
gt_depth = inputs["gt_depth"]
mask = inputs["val_mask"]
depth_errors = compute_depth_metrics(gt_depth, pred_depth,
mask)
for i, key in enumerate(self.evaluator.metrics.keys()):
losses[key] = np.array(depth_errors[i].cpu())
self.log("train", inputs, outputs, losses)
self.step += 1
def process_batch(self, inputs):
for key, ipt in inputs.items():
if key not in ["rgb", "cube_rgb"]:
inputs[key] = ipt.to(self.device)
losses = {}
equi_inputs = inputs["normalized_rgb"]
cube_inputs = inputs["normalized_cube_rgb"]
outputs = self.model(equi_inputs, cube_inputs)
losses["loss"] = self.compute_loss(inputs["gt_depth"],
outputs["pred_depth"],
inputs["val_mask"])
return outputs, losses
def validate(self):
"""Validate the model on the validation set
"""
self.model.eval()
self.evaluator.reset_eval_metrics()
pbar = tqdm.tqdm(self.val_loader)
pbar.set_description("Validating Epoch_{}".format(self.epoch))
with torch.no_grad():
for batch_idx, inputs in enumerate(pbar):
outputs, losses = self.process_batch(inputs)
pred_depth = outputs["pred_depth"].detach()
gt_depth = inputs["gt_depth"]
mask = inputs["val_mask"]
self.evaluator.compute_eval_metrics(gt_depth, pred_depth, mask)
for i, key in enumerate(self.evaluator.metrics.keys()):
losses[key] = np.array(self.evaluator.metrics[key].avg.cpu())
self.log("val", inputs, outputs, losses)
del inputs, outputs, losses
def log(self, mode, inputs, outputs, losses):
"""Write an event to the tensorboard events file
"""
writer = self.writers[mode]
for l, v in losses.items():
writer.add_scalar("{}".format(l), v, self.step)
for j in range(min(
4,
self.settings.batch_size)): # write a maxmimum of four images
writer.add_image("rgb/{}".format(j), inputs["rgb"][j].data,
self.step)
writer.add_image("cube_rgb/{}".format(j),
inputs["cube_rgb"][j].data, self.step)
writer.add_image(
"gt_depth/{}".format(j),
inputs["gt_depth"][j].data / inputs["gt_depth"][j].data.max(),
self.step)
writer.add_image(
"pred_depth/{}".format(j), outputs["pred_depth"][j].data /
outputs["pred_depth"][j].data.max(), self.step)
def save_settings(self):
"""Save settings to disk so we know what we ran this experiment with
"""
models_dir = os.path.join(self.log_path, "models")
if not os.path.exists(models_dir):
os.makedirs(models_dir)
to_save = self.settings.__dict__.copy()
with open(os.path.join(models_dir, 'settings.json'), 'w') as f:
json.dump(to_save, f, indent=2)
def save_model(self):
"""Save model weights to disk
"""
save_folder = os.path.join(self.log_path, "models",
"weights_{}".format(self.epoch))
if not os.path.exists(save_folder):
os.makedirs(save_folder)
save_path = os.path.join(save_folder, "{}.pth".format("model"))
to_save = self.model.state_dict()
# save resnet layers - these are needed at prediction time
to_save['layers'] = self.settings.num_layers
# save the input sizes
to_save['height'] = self.settings.height
to_save['width'] = self.settings.width
# save the dataset to train on
to_save['dataset'] = self.settings.dataset
to_save['net'] = self.settings.net
to_save['fusion'] = self.settings.fusion
to_save['se_in_fusion'] = self.settings.se_in_fusion
torch.save(to_save, save_path)
save_path = os.path.join(save_folder, "{}.pth".format("adam"))
torch.save(self.optimizer.state_dict(), save_path)
def load_model(self):
"""Load model from disk
"""
self.settings.load_weights_dir = os.path.expanduser(
self.settings.load_weights_dir)
assert os.path.isdir(self.settings.load_weights_dir), \
"Cannot find folder {}".format(self.settings.load_weights_dir)
print("loading model from folder {}".format(
self.settings.load_weights_dir))
path = os.path.join(self.settings.load_weights_dir,
"{}.pth".format("model"))
model_dict = self.model.state_dict()
pretrained_dict = torch.load(path)
pretrained_dict = {
k: v
for k, v in pretrained_dict.items() if k in model_dict
}
model_dict.update(pretrained_dict)
self.model.load_state_dict(model_dict)
# loading adam state
optimizer_load_path = os.path.join(self.settings.load_weights_dir,
"adam.pth")
if os.path.isfile(optimizer_load_path):
print("Loading Adam weights")
optimizer_dict = torch.load(optimizer_load_path)
self.optimizer.load_state_dict(optimizer_dict)
else:
print("Cannot find Adam weights so Adam is randomly initialized")
def __init__(self,
env_cls,
args,
manager,
cfg,
process_num,
pre=False,
pre_irl=False,
infer=False,
realtrain=False):
"""
:param env_cls: env class or function, not instance, as we need to create several instance in class.
:param args:
:param manager:
:param cfg:
:param process_num: process number
:param pre: set to pretrain mode
:param infer: set to test mode
"""
self.cfg = cfg
self.policy_clip = args.policy_clip
self.train_direc = args.train_direc
self.realtrain = realtrain
self.process_num = process_num
self.human_reward = human_reward(args, cfg)
self.gan_type = args.gan_type
# initialize envs for each process
self.env_list = []
for _ in range(process_num):
self.env_list.append(env_cls())
self.policy = MultiDiscretePolicy(cfg).to(device=DEVICE)
self.value = Value(cfg).to(device=DEVICE)
logging.info(summary(self.policy, show_weights=False))
# logging.info(summary(self.value , show_weights=False))
self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
self.loss = nn.BCELoss()
# if pre:
self.print_per_batch = args.print_per_batch
from dbquery import DBQuery
db = DBQuery(args.data_dir)
self.data_train = manager.create_dataset_seq('train', args.batchsz,
cfg, db)
self.data_valid = manager.create_dataset_seq('valid', args.batchsz,
cfg, db)
self.data_test = manager.create_dataset_seq('test', args.batchsz, cfg,
db)
self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
self.evaluator = Evaluator(args.data_dir, cfg)
# else:
# # self.rewarder = RewardEstimator(args, manager, cfg, pretrain=pre_irl, inference=infer, feature_extractor=self.policy)
# self.rewarder = DiscEstimator(args, manager, cfg, pretrain=pre_irl, inference=infer)
# self.evaluator = Evaluator(args.data_dir, cfg)
# from dbquery import DBQuery
# db = DBQuery(args.data_dir)
# self.data_train = manager.create_dataset_seq('train', args.batchsz, cfg, db)
self.expert_iter = iter(self.data_train)
self.save_dir = args.save_dir
self.save_per_epoch = args.save_per_epoch
self.optim_batchsz = args.batchsz
self.update_round = args.update_round
self.policy.eval()
self.gamma = args.gamma
self.epsilon = args.epsilon
self.tau = args.tau
self.policy_optim = optim.Adam(self.policy.parameters(), lr=args.lr_rl)
self.value_optim = optim.Adam(self.value.parameters(), lr=args.lr_rl)
self.mt_factor = args.mt_factor
self.epoch = 0
self.best_valid_loss = np.inf
self.valid_loss_threshold = np.inf
self.patience = 10
class DiaSeq(object):
def __init__(self,
env_cls,
args,
manager,
cfg,
process_num,
pre=False,
pre_irl=False,
infer=False,
realtrain=False):
"""
:param env_cls: env class or function, not instance, as we need to create several instance in class.
:param args:
:param manager:
:param cfg:
:param process_num: process number
:param pre: set to pretrain mode
:param infer: set to test mode
"""
self.cfg = cfg
self.policy_clip = args.policy_clip
self.train_direc = args.train_direc
self.realtrain = realtrain
self.process_num = process_num
self.human_reward = human_reward(args, cfg)
self.gan_type = args.gan_type
# initialize envs for each process
self.env_list = []
for _ in range(process_num):
self.env_list.append(env_cls())
self.policy = MultiDiscretePolicy(cfg).to(device=DEVICE)
self.value = Value(cfg).to(device=DEVICE)
logging.info(summary(self.policy, show_weights=False))
# logging.info(summary(self.value , show_weights=False))
self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
self.loss = nn.BCELoss()
# if pre:
self.print_per_batch = args.print_per_batch
from dbquery import DBQuery
db = DBQuery(args.data_dir)
self.data_train = manager.create_dataset_seq('train', args.batchsz,
cfg, db)
self.data_valid = manager.create_dataset_seq('valid', args.batchsz,
cfg, db)
self.data_test = manager.create_dataset_seq('test', args.batchsz, cfg,
db)
self.multi_entropy_loss = nn.MultiLabelSoftMarginLoss()
self.evaluator = Evaluator(args.data_dir, cfg)
# else:
# # self.rewarder = RewardEstimator(args, manager, cfg, pretrain=pre_irl, inference=infer, feature_extractor=self.policy)
# self.rewarder = DiscEstimator(args, manager, cfg, pretrain=pre_irl, inference=infer)
# self.evaluator = Evaluator(args.data_dir, cfg)
# from dbquery import DBQuery
# db = DBQuery(args.data_dir)
# self.data_train = manager.create_dataset_seq('train', args.batchsz, cfg, db)
self.expert_iter = iter(self.data_train)
self.save_dir = args.save_dir
self.save_per_epoch = args.save_per_epoch
self.optim_batchsz = args.batchsz
self.update_round = args.update_round
self.policy.eval()
self.gamma = args.gamma
self.epsilon = args.epsilon
self.tau = args.tau
self.policy_optim = optim.Adam(self.policy.parameters(), lr=args.lr_rl)
self.value_optim = optim.Adam(self.value.parameters(), lr=args.lr_rl)
self.mt_factor = args.mt_factor
self.epoch = 0
self.best_valid_loss = np.inf
self.valid_loss_threshold = np.inf
self.patience = 10
def retrieve_expert(self, to_device_=True):
try:
data = self.expert_iter.next()
except StopIteration:
self.expert_iter = iter(self.data_train)
data = self.expert_iter.next()
self.epoch += 1
if to_device_:
return to_device(data)
else:
return data
def policy_loop(self, data):
s, target_a, target_d = to_device(data)
target_a = id2onehot(target_a)
a_weights = self.policy(s)
loss_a = self.multi_entropy_loss(a_weights, target_a)
return loss_a
def imitate_loop(self, data):
s, target_a, target_d = to_device(data)
target_a = id2onehot(target_a)
a_weights = self.policy.imitate(s)
loss_a = self.multi_entropy_loss(a_weights, target_a)
return loss_a
def prepare_data(self, data_batch):
s_real, a_real_, next_s_real, d_real = to_device(data_batch)
def update_loop(self, data):
# gen_type: beam, greedy
# mode: GEN, TEACH_FORCE
# s, a, d, a_seq = self.retrieve_expert(True)
s, _, _, a_seq = data
s = s.to(device=DEVICE)
a_seq = a_seq.to(device=DEVICE)
loss = self.policy(s=s,
a_seq=a_seq[:, :self.cfg.max_len],
mode=TEACH_FORCE)
return loss
def train(self, epoch):
"""
pretrain the policy by simple imitation learning (behavioral cloning)
"""
self.policy.train()
a_loss = 0.
for i, data in enumerate(self.data_train):
self.policy_optim.zero_grad()
loss_a = self.update_loop(data)
a_loss += loss_a.item()
loss_a.backward()
torch.nn.utils.clip_grad_norm_(self.policy.parameters(),
self.policy_clip)
self.policy_optim.step()
if (i + 1) % self.print_per_batch == 0:
a_loss /= self.print_per_batch
logging.debug(
'<<dialog policy>> epoch {}, iter {}, loss_a:{}'.format(
epoch, i, a_loss))
a_loss = 0.
self.policy.eval()
if (epoch + 1) % self.save_per_epoch == 0:
self.save(self.save_dir, epoch, True)
for i, data in enumerate(self.data_valid):
loss_a = self.update_loop(data)
a_loss += loss_a.item()
valid_loss = a_loss / len(self.data_valid)
logging.debug(
'<<dialog policy>> validation, epoch {}, loss_a:{}'.format(
epoch, valid_loss))
if valid_loss < self.best_valid_loss:
if valid_loss <= self.valid_loss_threshold * self.cfg.improve_threshold:
self.patience = max(self.patience,
epoch * self.cfg.patient_increase)
self.valid_loss_threshold = valid_loss
logging.info("Update patience to {}".format(self.patience))
self.best_valid_loss = valid_loss
logging.info('<<dialog policy>> best model saved')
self.save(self.save_dir, 'best', True)
if self.cfg.early_stop and self.patience <= epoch:
if epoch < self.cfg.max_epoch:
logging.info("!!Early stop due to run out of patience!!")
logging.info("Best validation loss %f" % self.best_valid_loss)
return True
return False
def sample(self, batchsz):
"""
Given batchsz number of task, the batchsz will be splited equally to each processes
and when processes return, it merge all data and return
:param batchsz:
:return: batch
"""
# batchsz will be splitted into each process,
# final batchsz maybe larger than batchsz parameters
process_batchsz = np.ceil(batchsz / self.process_num).astype(np.int32)
# buffer to save all data
queue = mp.Queue()
# start processes for pid in range(1, processnum)
# if processnum = 1, this part will be ignored.
# when save tensor in Queue, the process should keep alive till Queue.get(),
# please refer to : https://discuss.pytorch.org/t/using-torch-tensor-over-multiprocessing-queue-process-fails/2847
# however still some problem on CUDA tensors on multiprocessing queue,
# please refer to : https://discuss.pytorch.org/t/cuda-tensors-on-multiprocessing-queue/28626
# so just transform tensors into numpy, then put them into queue.
evt = mp.Event()
processes = []
for i in range(self.process_num):
process_args = (i, queue, evt, self.env_list[i], self.policy,
process_batchsz, self.human_reward)
processes.append(mp.Process(target=sampler, args=process_args))
for p in processes:
# set the process as daemon, and it will be killed once the main process is stoped.
p.daemon = True
p.start()
# we need to get the first Memory object and then merge others Memory use its append function.
pid0, buff0 = queue.get()
for _ in range(1, self.process_num):
pid, buff_ = queue.get()
buff0.append(buff_) # merge current Memory into buff0
evt.set()
# now buff saves all the sampled data
buff = buff0
return buff.get_batch()
def evaluate(self, save_dialog=False):
self.policy.eval()
if save_dialog:
with open('./data/goal.json', 'r') as f:
saved_goal_list = json.load(f)
collected_dialog = []
env = self.env_list[0]
traj_len = 40
reward_tot, turn_tot, inform_tot, match_tot, success_tot = [], [], [], [], []
for seed in range(1000):
dialog_list = []
if save_dialog:
s = env.reset(seed, saved_goal=saved_goal_list[seed])
else:
s = env.reset(seed, saved_goal=None)
# print('seed', seed)
# print('goal', env.goal.domain_goals)
# print('usr', s['user_action'])
dialog_list.append(s['user_action'])
turn = traj_len
reward = []
value = []
mask = []
for t in range(traj_len):
s_vec = torch.Tensor(state_vectorize(s, env.cfg,
env.db)).to(device=DEVICE)
# mode with policy during evaluation
a = self.policy.select_action(s_vec, False)
next_s, done = env.step(s, a.cpu())
next_s_vec = torch.Tensor(
state_vectorize(next_s, env.cfg, env.db)).to(device=DEVICE)
r = self.human_reward.reward_human(next_s, done)
reward.append(r)
s = next_s
dialog_list.append(s['last_sys_action'])
dialog_list.append(s['user_action'])
# print('sys', s['last_sys_action'])
# print('usr', s['user_action'])
if done:
mask.append(0)
turn = t + 2 # one due to counting from 0, the one for the last turn
break
mask.append(1)
reward_tot.append(np.mean(reward))
turn_tot.append(turn)
match_tot += self.evaluator.match_rate(s)
inform_tot.append(self.evaluator.inform_F1(s))
reward = torch.Tensor(reward)
mask = torch.LongTensor(mask)
# print('turn', turn)
match_session = self.evaluator.match_rate(s, True)
# print('match', match_session)
inform_session = self.evaluator.inform_F1(s, True)
# print('inform', inform_session)
if (match_session == 1 and inform_session[1] == 1) \
or (match_session == 1 and inform_session[1] is None) \
or (match_session is None and inform_session[1] == 1):
# print('success', 1)
success_tot.append(1)
else:
# print('success', 0)
success_tot.append(0)
dialog_dict = {
'goal id': seed,
'goal': env.goal.domain_goals,
'dialog': dialog_list,
'turn': turn,
'status': success_tot[-1]
}
collected_dialog.append(dialog_dict)
logging.info('reward {}'.format(np.mean(reward_tot)))
logging.info('turn {}'.format(np.mean(turn_tot)))
logging.info('match {}'.format(np.mean(match_tot)))
TP, FP, FN = np.sum(inform_tot, 0)
prec = TP / (TP + FP)
rec = TP / (TP + FN)
F1 = 2 * prec * rec / (prec + rec)
logging.info('inform rec {}, F1 {}'.format(rec, F1))
logging.info('success {}'.format(np.mean(success_tot)))
if save_dialog:
self.save_dialog(self.save_dir, collected_dialog)
def save_dialog(self, directory, collected_dialog):
if not os.path.exists(directory):
os.makedirs(directory)
des_path = directory + '/' + 'collected_dialog.json'
with open(des_path, 'w') as f:
json.dump(collected_dialog, f, indent=4)
def save(self, directory, epoch, rl_only=False):
if not os.path.exists(directory):
os.makedirs(directory)
# if not rl_only:
# self.rewarder.save_irl(directory, epoch)
torch.save(self.value.state_dict(),
directory + '/' + str(epoch) + '_ppo.val.mdl')
torch.save(self.policy.state_dict(),
directory + '/' + str(epoch) + '_ppo.pol.mdl')
logging.info(
'<<dialog policy>> epoch {}: saved network to mdl'.format(epoch))
def load(self, filename):
# self.rewarder.load_irl(filename)
# if self.realtrain:
# filename='model_saved_ori/model_agenda_pre_ori/best'
# filename='model_saved/model_agenda_pre_mt_op_1.0/best'
value_mdl = filename + '_ppo.val.mdl'
policy_mdl = filename + '_ppo.pol.mdl'
if os.path.exists(value_mdl):
self.value.load_state_dict(torch.load(value_mdl))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
value_mdl))
if os.path.exists(policy_mdl):
self.policy.load_state_dict(torch.load(policy_mdl))
logging.info(
'<<dialog policy>> loaded checkpoint from file: {}'.format(
policy_mdl))
best_pkl = filename + '.pkl'
if os.path.exists(best_pkl):
with open(best_pkl, 'rb') as f:
best = pickle.load(f)
else:
best = [float('inf'), float('inf'), float('-inf')]
return best
def __init__(self, settings):
self.settings = settings
self.device = torch.device(
"cuda" if len(self.settings.gpu_devices) else "cpu")
self.gpu_devices = ','.join([str(id) for id in settings.gpu_devices])
os.environ["CUDA_VISIBLE_DEVICES"] = self.gpu_devices
self.log_path = os.path.join(self.settings.log_dir,
self.settings.model_name)
# checking the input height and width are multiples of 32
assert self.settings.height % 32 == 0, "input height must be a multiple of 32"
assert self.settings.width % 32 == 0, "input width must be a multiple of 32"
# data
datasets_dict = {
"3d60": datasets.ThreeD60,
"panosuncg": datasets.PanoSunCG,
"stanford2d3d": datasets.Stanford2D3D,
"matterport3d": datasets.Matterport3D
}
self.dataset = datasets_dict[self.settings.dataset]
self.settings.cube_w = self.settings.height // 2
fpath = os.path.join(os.path.dirname(__file__), "datasets",
"{}_{}.txt")
train_file_list = fpath.format(self.settings.dataset, "train")
val_file_list = fpath.format(self.settings.dataset, "val")
train_dataset = self.dataset(
self.settings.data_path,
train_file_list,
self.settings.height,
self.settings.width,
self.settings.disable_color_augmentation,
self.settings.disable_LR_filp_augmentation,
self.settings.disable_yaw_rotation_augmentation,
is_training=True)
self.train_loader = DataLoader(train_dataset,
self.settings.batch_size,
True,
num_workers=self.settings.num_workers,
pin_memory=True,
drop_last=True)
num_train_samples = len(train_dataset)
self.num_total_steps = num_train_samples // self.settings.batch_size * self.settings.num_epochs
val_dataset = self.dataset(
self.settings.data_path,
val_file_list,
self.settings.height,
self.settings.width,
self.settings.disable_color_augmentation,
self.settings.disable_LR_filp_augmentation,
self.settings.disable_yaw_rotation_augmentation,
is_training=False)
self.val_loader = DataLoader(val_dataset,
self.settings.batch_size,
False,
num_workers=self.settings.num_workers,
pin_memory=True,
drop_last=True)
# network
Net_dict = {"UniFuse": UniFuse, "Equi": Equi}
Net = Net_dict[self.settings.net]
self.model = Net(self.settings.num_layers,
self.settings.height,
self.settings.width,
self.settings.imagenet_pretrained,
train_dataset.max_depth_meters,
fusion_type=self.settings.fusion,
se_in_fusion=self.settings.se_in_fusion)
self.model.to(self.device)
self.parameters_to_train = list(self.model.parameters())
self.optimizer = optim.Adam(self.parameters_to_train,
self.settings.learning_rate)
if self.settings.load_weights_dir is not None:
self.load_model()
print("Training model named:\n ", self.settings.model_name)
print("Models and tensorboard events files are saved to:\n",
self.settings.log_dir)
print("Training is using:\n ", self.device)
self.compute_loss = BerhuLoss()
self.evaluator = Evaluator()
self.writers = {}
for mode in ["train", "val"]:
self.writers[mode] = SummaryWriter(
os.path.join(self.log_path, mode))
self.save_settings()
def __init__(self, args, config):
self.evaluator = Evaluator(args.data_dir, config)
def model_fn(features, labels, mode, params, config):
"""This is a function for creating a computational tensorflow graph.
The function is in format required by tf.estimator.
"""
# choose a backbone network
if params['backbone'] == 'resnet':
feature_extractor = resnet
checkpoint_scope = 'resnet_v1_50/'
elif params['backbone'] == 'mobilenet':
feature_extractor = lambda x: mobilenet(x, params['depth_multiplier'])
checkpoint_scope = 'MobilenetV2/'
elif params['backbone'] == 'shufflenet':
feature_extractor = lambda x: shufflenet(
x, str(params['depth_multiplier']))
checkpoint_scope = 'ShuffleNetV2/'
# build the main graph
is_training = mode == tf.estimator.ModeKeys.TRAIN
detector = Detector(features['images'], feature_extractor, is_training,
params)
# use a pretrained backbone network
if is_training:
with tf.name_scope('init_from_checkpoint'):
tf.train.init_from_checkpoint(params['pretrained_checkpoint'],
{checkpoint_scope: checkpoint_scope})
# add NMS to the graph
if not is_training:
predictions = detector.get_predictions(
score_threshold=params['score_threshold'],
iou_threshold=params['iou_threshold'],
max_boxes_per_class=params['max_boxes_per_class'])
if mode == tf.estimator.ModeKeys.PREDICT:
w, h = tf.unstack(tf.to_float(
features['images_size'])) # original image size
s = tf.to_float(tf.shape(features['images'])) # size after resizing
scaler = tf.stack([h / s[1], w / s[2], h / s[1], w / s[2]])
predictions['boxes'] = scaler * predictions['boxes']
export_outputs = tf.estimator.export.PredictOutput({
name: tf.identity(tensor, name)
for name, tensor in predictions.items()
})
return tf.estimator.EstimatorSpec(
mode,
predictions=predictions,
export_outputs={'outputs': export_outputs})
# add L2 regularization
with tf.name_scope('weight_decay'):
add_weight_decay(params['weight_decay'])
regularization_loss = tf.losses.get_regularization_loss()
# create localization and classification losses
losses = detector.get_losses(labels, params)
tf.losses.add_loss(params['alpha'] * losses['rpn_localization_loss'])
tf.losses.add_loss(params['beta'] * losses['rpn_classification_loss'])
tf.losses.add_loss(params['gamma'] * losses['roi_localization_loss'])
tf.losses.add_loss(params['theta'] * losses['roi_classification_loss'])
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
tf.summary.scalar('regularization_loss', regularization_loss)
tf.summary.scalar('rpn_localization_loss', losses['rpn_localization_loss'])
tf.summary.scalar('rpn_classification_loss',
losses['rpn_classification_loss'])
tf.summary.scalar('roi_localization_loss', losses['roi_localization_loss'])
tf.summary.scalar('roi_classification_loss',
losses['roi_classification_loss'])
if mode == tf.estimator.ModeKeys.EVAL:
with tf.name_scope('evaluator'):
evaluator = Evaluator(num_classes=params['num_classes'])
eval_metric_ops = evaluator.get_metric_ops(labels, predictions)
return tf.estimator.EstimatorSpec(mode,
loss=total_loss,
eval_metric_ops=eval_metric_ops)
assert mode == tf.estimator.ModeKeys.TRAIN
with tf.variable_scope('learning_rate'):
global_step = tf.train.get_global_step()
learning_rate = tf.train.piecewise_constant(global_step,
params['lr_boundaries'],
params['lr_values'])
tf.summary.scalar('learning_rate', learning_rate)
with tf.variable_scope('optimizer'):
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=MOMENTUM,
use_nesterov=USE_NESTEROV)
grads_and_vars = optimizer.compute_gradients(total_loss)
grads_and_vars = [(tf.clip_by_norm(g, GRADIENT_CLIP), v)
for g, v in grads_and_vars]
train_op = optimizer.apply_gradients(grads_and_vars, global_step)
for g, v in grads_and_vars:
tf.summary.histogram(v.name[:-2] + '_hist', v)
tf.summary.histogram(v.name[:-2] + '_grad_hist', g)
with tf.control_dependencies([train_op]), tf.name_scope('ema'):
ema = tf.train.ExponentialMovingAverage(decay=MOVING_AVERAGE_DECAY,
num_updates=global_step)
train_op = ema.apply(tf.trainable_variables())
return tf.estimator.EstimatorSpec(mode, loss=total_loss, train_op=train_op)
def display_baseline(
config: Config,
graph: Graph,
train_trajectories: Trajectories,
test_trajectories: Trajectories,
evaluator: Evaluator,
):
"""Compute baseline uniform random walk with/without
Args:
config (Config): [description]
graph (Graph): graph
train_trajectories (Trajectories): train trajectories
test_trajectories (Trajectories): test trajectories
evaluator (Evaluator): evaluator to compute metrics
"""
graph = graph.add_self_loops(degree_zero_only=True)
graph = graph.update(edges=torch.ones(graph.n_edge, device=graph.device))
graph = graph.softmax_weights()
print("Computing non backtracking edges...")
graph.compute_non_backtracking_edges()
print("Done")
print("=== BASELINE ===")
baseline_model = create_baseline(config, non_backtracking=False)
print("TEST DATASET")
evaluator.compute(baseline_model, graph, test_trajectories, None)
print(colored(evaluator.to_string(), "green"))
print("TRAIN DATASET")
evaluator.compute(baseline_model, graph, train_trajectories, None)
print(colored(evaluator.to_string(), "green"))
print("=== NON BACKTRACKING BASELINE ===")
nb_baseline_model = create_baseline(config, non_backtracking=True)
print("TEST DATASET")
evaluator.compute(nb_baseline_model, graph, test_trajectories, None)
print(colored(evaluator.to_string(), "green"))
print("TRAIN DATASET")
evaluator.compute(nb_baseline_model, graph, train_trajectories, None)
print(colored(evaluator.to_string(), "green"))
def main():
# 配置设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 数据路径
train_dir = './data_road/training'
# 超参数
num_epochs = 60
learning_rate = 0.001
img_size = (640, 192)
num_class = 2
SAVE_INTERVAL = 5
evaluator = Evaluator(num_class=num_class)
# 构建Dataset实例
train_data = LaneDataset(data_dir=train_dir, img_size=img_size)
train_loader = DataLoader(train_data, batch_size=2, shuffle=True)
# 构建Model实例
model = FCN8s(n_class=2)
vgg16_path = "models/vgg16_from_caffe.pth"
vgg16 = load_vgg16(model_file=vgg16_path)
model.copy_params_from_vgg16(vgg16)
model = model.to(device)
print("模型加载成功！！！")
# 定义损失函数和优化器
criterion = nn.BCELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=5e-4)
Acc_list = []
Acc_class_list = []
mIoU_list = []
FWIoU_list = []
loss_list = []
# 训练
for epoch in range(num_epochs):
train_loss = 0.0
evaluator.reset()
for i, (image, label) in enumerate(train_loader):
image = image.to(device)
label = label.to(device)
optimizer.zero_grad()
# 前向传播
output = model(image)
output = torch.sigmoid(output)
loss = criterion(output, label)
output = torch.argmax(output, dim=1).cpu().numpy()
label = torch.argmax(label, dim=1).cpu().numpy()
# 后向传播及优化
loss.backward()
train_loss += loss.item()
optimizer.step()
evaluator.add_batch(label, output) # 添加output和label，用于后续评估
# 计算像素准确率、平均像素准确率、平均交并比、频权交并比
Acc = evaluator.Pixel_Accuracy()
Acc_class = evaluator.Pixel_Accuracy_Class()
mIoU = evaluator.Mean_Intersection_over_Union()
FWIoU = evaluator.Frequency_Weighted_Intersection_over_Union()
Acc_list.append(Acc)
Acc_class_list.append(Acc_class)
mIoU_list.append(mIoU)
FWIoU_list.append(FWIoU)
loss_list.append(train_loss/len(train_loader))
# 保存模型
if epoch % SAVE_INTERVAL == 0 or (epoch+1) == num_epochs:
torch.save(model.state_dict(), './models/fcn8s_{}.pth'.format(epoch))
print("Epoch_{}: train_loss: {:.6f} Acc: {:.4f} Acc_class: {:.4f} mIoU: {:.4f} FWIoU: {:.4f}"
.format(epoch+1, train_loss/len(train_loader), Acc, Acc_class, mIoU, FWIoU))
# 绘制曲线
draw_figures(loss_list, title='train_loss')
draw_figures(Acc_list, title='Acc')
draw_figures(Acc_class_list, title='Acc_class')
draw_figures(mIoU_list, title='mIoU')
draw_figures(FWIoU_list, title='FWIoU')
print("完成曲线绘制！！！")
def main():
# print("BEGIN...")
parser = argparse.ArgumentParser()
parser.add_argument("config_file")
parser.add_argument("--name")
parser.add_argument("--bs")
parser.add_argument("--lr")
parser.add_argument("--loss")
args = parser.parse_args()
# load configuration
writer = SummaryWriter()
config = Config()
config.load_from_file(args.config_file)
# Hyperparameter logging
if 'bs' in args:
config.batch_size = int(args.bs)
if 'lr' in args:
config.lr = float(args.lr)
if 'loss' in args:
config.loss = args.loss
writer.add_scalar('batch_size', config.batch_size)
writer.add_scalar('lr', config.lr)
writer.add_text('loss_func', config.loss)
graph, trajectories, pairwise_node_features, _ = load_data(config)
print("loaded data...")
train_trajectories, valid_trajectories, test_trajectories = trajectories
use_validation_set = len(valid_trajectories) > 0
graph = graph.to(config.device)
print("loaded device...")
given_as_target, siblings_nodes = None, None
if config.dataset == "wikispeedia":
given_as_target, siblings_nodes = load_wiki_data(config)
if pairwise_node_features is not None:
pairwise_node_features = pairwise_node_features.to(config.device)
if config.rw_edge_weight_see_number_step: # TODO
raise NotImplementedError
if args.name is not None:
print(f"Experiment name from CLI: {args.name}")
config.name = args.name
if not config.name:
experiment_name = input("Give a name to this experiment? ").strip()
config.name = experiment_name or config.date
print(f'==== START "{config.name}" ====')
torch.manual_seed(config.seed)
if config.enable_checkpointing:
chkpt_dir = os.path.join(config.workspace, config.checkpoint_directory,
config.name)
os.makedirs(chkpt_dir, exist_ok=True)
print(f"Checkpoints will be saved in [{chkpt_dir}]")
d_node = graph.nodes.shape[1] if graph.nodes is not None else 0
d_edge = graph.edges.shape[1] if graph.edges is not None else 0
print(
f"Number of node features {d_node}. Number of edge features {d_edge}")
model = create_model(graph, pairwise_node_features, config)
model = model.to(config.device)
optimizer = create_optimizer(model.parameters(), config)
if config.restore_from_checkpoint:
filename = input("Checkpoint file: ")
checkpoint_data = torch.load(filename)
model.load_state_dict(checkpoint_data["model_state_dict"])
optimizer.load_state_dict(checkpoint_data["optimizer_state_dict"])
print("Loaded parameters from checkpoint")
def create_evaluator():
return Evaluator(
graph.n_node,
given_as_target=given_as_target,
siblings_nodes=siblings_nodes,
config=config,
)
if use_validation_set:
valid_evaluator = Evaluator(
graph.n_node,
given_as_target=given_as_target,
siblings_nodes=siblings_nodes,
config=config,
)
if config.compute_baseline:
display_baseline(config, graph, train_trajectories, test_trajectories,
create_evaluator())
graph = graph.add_self_loops(
degree_zero_only=config.self_loop_deadend_only,
edge_value=config.self_loop_weight)
if config.rw_non_backtracking:
print("Computing non backtracking graph...", end=" ")
sys.stdout.flush()
graph.compute_non_backtracking_edges()
print("Done")
evaluate(model,
graph,
test_trajectories,
pairwise_node_features,
create_evaluator,
dataset="TEST")
if use_validation_set:
evaluate(
model,
graph,
valid_trajectories,
pairwise_node_features,
create_evaluator,
dataset="EVAL",
)
for epoch in range(config.number_epoch):
print(f"\n=== EPOCH {epoch} ===")
model.train()
train_epoch(model, graph, optimizer, config, train_trajectories,
pairwise_node_features)
# VALID and TEST metrics computation
test_evaluator = evaluate(
model,
graph,
test_trajectories,
pairwise_node_features,
create_evaluator,
dataset="TEST",
)
writer.add_scalar('precision_top1',
test_evaluator.metrics['precision_top1'].mean(),
epoch)
writer.add_scalar('precision_top5',
test_evaluator.metrics['precision_top5'].mean(),
epoch)
writer.add_scalar('target_p',
test_evaluator.metrics['target_probability'].mean(),
epoch)
writer.add_scalar('choice_acc',
test_evaluator.metrics['choice_accuracy'].mean(),
epoch)
valid_evaluator = None
if use_validation_set:
valid_evaluator = evaluate(
model,
graph,
valid_trajectories,
pairwise_node_features,
create_evaluator,
dataset="EVAL",
)
if config.enable_checkpointing and epoch % config.chechpoint_every_num_epoch == 0:
print("Checkpointing...")
directory = os.path.join(config.workspace,
config.checkpoint_directory, config.name)
chkpt_file = os.path.join(directory, f"{epoch:04d}.pt")
torch.save(
{
"epoch": epoch,
"model_state_dict": model.state_dict(),
"optimizer_state_dict": optimizer.state_dict(),
},
chkpt_file,
)
config_file = os.path.join(directory, "config")
config.save_to_file(config_file)
metrics_file = os.path.join(directory, f"{epoch:04d}.txt")
with open(metrics_file, "w") as f:
f.write(test_evaluator.to_string())
if valid_evaluator:
f.write("\n\n=== VALIDATION ==\n\n")
f.write(valid_evaluator.to_string())
print(colored(f"Checkpoint saved in {chkpt_file}", "blue"))
