class ExponentialLRScheduler(Callback):
def __init__(self, gamma, epoch_every=1, batch_every=None):
super().__init__()
self.gamma = gamma
if epoch_every == 0:
self.epoch_every = False
else:
self.epoch_every = epoch_every
if batch_every == 0:
self.batch_every = False
else:
self.batch_every = batch_every
def set_params(self, transformer, validation_datagen, *args, **kwargs):
self.validation_datagen = validation_datagen
self.model = transformer.model
self.optimizer = transformer.optimizer
self.loss_function = transformer.loss_function
self.lr_scheduler = ExponentialLR(self.optimizer, self.gamma, last_epoch=-1)
def on_train_begin(self, *args, **kwargs):
self.epoch_id = 0
self.batch_id = 0
logger.info('initial lr: {0}'.format(self.optimizer.state_dict()['param_groups'][0]['initial_lr']))
def on_epoch_end(self, *args, **kwargs):
if self.epoch_every and (((self.epoch_id + 1) % self.epoch_every) == 0):
self.lr_scheduler.step()
logger.info('epoch {0} current lr: {1}'.format(self.epoch_id + 1,
self.optimizer.state_dict()['param_groups'][0]['lr']))
self.epoch_id += 1
def on_batch_end(self, *args, **kwargs):
if self.batch_every and ((self.batch_id % self.batch_every) == 0):
self.lr_scheduler.step()
logger.info('epoch {0} batch {1} current lr: {2}'.format(
self.epoch_id + 1, self.batch_id + 1, self.optimizer.state_dict()['param_groups'][0]['lr']))
self.batch_id += 1
class TrackerSiamFC(Tracker):
def __init__(self, net_path=None, **kwargs):
super(TrackerSiamFC, self).__init__('SiamFC', True)
self.cfg = self.parse_args(**kwargs)
# setup GPU device if available
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
# setup model
# self.net = Net(
# backbone=AlexNetV1(),
# head=SiamFC(self.cfg.out_scale)
# )
self.net = Net(backbone=CapsuleNetV1(),
head=SiamFC_V2(self.cfg.out_scale))
ops.init_weights(self.net)
# load checkpoint if provided
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
self.net = self.net.to(self.device)
# setup criterion
self.criterion = BalancedLoss()
# setup optimizer
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.cfg.initial_lr,
weight_decay=self.cfg.weight_decay,
momentum=self.cfg.momentum)
# setup lr scheduler
gamma = np.power(self.cfg.ultimate_lr / self.cfg.initial_lr,
1.0 / self.cfg.epoch_num)
self.lr_scheduler = ExponentialLR(self.optimizer, gamma)
def parse_args(self, **kwargs):
# default parameters
cfg = {
# basic parameters
'out_scale': 0.001,
'exemplar_sz': 127,
'instance_sz': 255,
'context': 0.5,
# inference parameters
'scale_num': 3,
'scale_step': 1.0375,
'scale_lr': 0.59,
'scale_penalty': 0.9745,
'window_influence': 0.176,
'response_sz': 17,
'response_up': 16,
'total_stride': 8,
# train parameters
'epoch_num': 50,
'batch_size': 8,
'num_workers': 32,
'initial_lr': 1e-2,
'ultimate_lr': 1e-5,
'weight_decay': 5e-4,
'momentum': 0.9,
'r_pos': 16,
'r_neg': 0
}
for key, val in kwargs.items():
if key in cfg:
cfg.update({key: val})
return namedtuple('Config', cfg.keys())(**cfg)
@torch.no_grad()
def init(self, img, box):
# set to evaluation mode
self.net.eval()
# convert box to 0-indexed and center based [y, x, h, w]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
self.center, self.target_sz = box[:2], box[2:]
# create hanning window
self.upscale_sz = self.cfg.response_up * self.cfg.response_sz
self.hann_window = np.outer(np.hanning(self.upscale_sz),
np.hanning(self.upscale_sz))
self.hann_window /= self.hann_window.sum()
# search scale factors
self.scale_factors = self.cfg.scale_step**np.linspace(
-(self.cfg.scale_num // 2), self.cfg.scale_num // 2,
self.cfg.scale_num)
# exemplar and search sizes
context = self.cfg.context * np.sum(self.target_sz)
self.z_sz = np.sqrt(np.prod(self.target_sz + context))
self.x_sz = self.z_sz * \
self.cfg.instance_sz / self.cfg.exemplar_sz
# exemplar image
self.avg_color = np.mean(img, axis=(0, 1))
z = ops.crop_and_resize(img,
self.center,
self.z_sz,
out_size=self.cfg.exemplar_sz,
border_value=self.avg_color)
# exemplar features
z = torch.from_numpy(z).to(self.device).permute(
2, 0, 1).unsqueeze(0).float()
self.kernel = self.net.backbone(z)
@torch.no_grad()
def update(self, img):
# set to evaluation mode
self.net.eval()
# search images
x = [
ops.crop_and_resize(img,
self.center,
self.x_sz * f,
out_size=self.cfg.instance_sz,
border_value=self.avg_color)
for f in self.scale_factors
]
x = np.stack(x, axis=0)
x = torch.from_numpy(x).to(self.device).permute(0, 3, 1, 2).float()
# responses
x = self.net.backbone(x)
responses = self.net.head(self.kernel, x)
responses = responses.squeeze(1).cpu().numpy()
# upsample responses and penalize scale changes
responses = np.stack([
cv2.resize(u, (self.upscale_sz, self.upscale_sz),
interpolation=cv2.INTER_CUBIC) for u in responses
])
responses[:self.cfg.scale_num // 2] *= self.cfg.scale_penalty
responses[self.cfg.scale_num // 2 + 1:] *= self.cfg.scale_penalty
# peak scale
scale_id = np.argmax(np.amax(responses, axis=(1, 2)))
# peak location
response = responses[scale_id]
response -= response.min()
response /= response.sum() + 1e-16
response = (1 - self.cfg.window_influence) * response + \
self.cfg.window_influence * self.hann_window
loc = np.unravel_index(response.argmax(), response.shape)
# locate target center
disp_in_response = np.array(loc) - (self.upscale_sz - 1) / 2
disp_in_instance = disp_in_response * \
self.cfg.total_stride / self.cfg.response_up
disp_in_image = disp_in_instance * self.x_sz * \
self.scale_factors[scale_id] / self.cfg.instance_sz
self.center += disp_in_image
# update target size
scale = (1 - self.cfg.scale_lr) * 1.0 + \
self.cfg.scale_lr * self.scale_factors[scale_id]
self.target_sz *= scale
self.z_sz *= scale
self.x_sz *= scale
# return 1-indexed and left-top based bounding box
box = np.array([
self.center[1] + 1 - (self.target_sz[1] - 1) / 2,
self.center[0] + 1 - (self.target_sz[0] - 1) / 2,
self.target_sz[1], self.target_sz[0]
])
return box
def track(self, img_files, box, visualize=False):
frame_num = len(img_files)
boxes = np.zeros((frame_num, 4))
boxes[0] = box
times = np.zeros(frame_num)
for f, img_file in enumerate(img_files):
img = ops.read_image(img_file)
begin = time.time()
if f == 0:
self.init(img, box)
else:
boxes[f, :] = self.update(img)
times[f] = time.time() - begin
if visualize:
ops.show_image(img, boxes[f, :])
return boxes, times
def train_step(self, batch, backward=True):
# set network mode
self.net.train(backward)
# parse batch data
z = batch[0].to(self.device, non_blocking=self.cuda)
x = batch[1].to(self.device, non_blocking=self.cuda)
with torch.set_grad_enabled(backward):
# inference
responses = self.net(z, x)
# calculate loss
labels = self._create_labels(responses.size())
loss = self.criterion(responses, labels)
if backward:
# back propagation
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
@torch.enable_grad()
def train_over(self, seqs, val_seqs=None, save_dir='pretrained'):
# set to train mode
self.net.train()
# create save_dir folder
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# setup dataset
transforms = SiamFCTransforms(exemplar_sz=self.cfg.exemplar_sz,
instance_sz=self.cfg.instance_sz,
context=self.cfg.context)
dataset = Pair(seqs=seqs, transforms=transforms)
# setup dataloader
dataloader = DataLoader(dataset,
batch_size=self.cfg.batch_size,
shuffle=True,
num_workers=self.cfg.num_workers,
pin_memory=self.cuda,
drop_last=True)
# loop over epochs
for epoch in range(self.cfg.epoch_num):
# update lr at each epoch
self.lr_scheduler.step(epoch=epoch)
# loop over dataloader
for it, batch in enumerate(dataloader):
loss = self.train_step(batch, backward=True)
print('Epoch: {} [{}/{}] Loss: {:.5f}'.format(
epoch + 1, it + 1, len(dataloader), loss))
sys.stdout.flush()
# save checkpoint
if not os.path.exists(save_dir):
os.makedirs(save_dir)
net_path = os.path.join(save_dir,
'siamfc_alexnet_e%d.pth' % (epoch + 1))
torch.save(self.net.state_dict(), net_path)
def _create_labels(self, size):
# skip if same sized labels already created
if hasattr(self, 'labels') and self.labels.size() == size:
return self.labels
def logistic_labels(x, y, r_pos, r_neg):
dist = np.abs(x) + np.abs(y) # block distance
labels = np.where(
dist <= r_pos, np.ones_like(x),
np.where(dist < r_neg,
np.ones_like(x) * 0.5, np.zeros_like(x)))
return labels
# distances along x- and y-axis
n, c, h, w = size
x = np.arange(w) - (w - 1) / 2
y = np.arange(h) - (h - 1) / 2
x, y = np.meshgrid(x, y)
# create logistic labels
r_pos = self.cfg.r_pos / self.cfg.total_stride
r_neg = self.cfg.r_neg / self.cfg.total_stride
labels = logistic_labels(x, y, r_pos, r_neg)
# repeat to size
labels = labels.reshape((1, 1, h, w))
labels = np.tile(labels, (n, c, 1, 1))
# convert to tensors
self.labels = torch.from_numpy(labels).to(self.device).float()
return self.labels
def train_and_eval(self):
torch.set_num_threads(2)
best_valid = [0, 0, 0, 0, 0]
best_test = [0, 0, 0, 0, 0]
self.entity_idxs = {d.entities[i]: i for i in range(len(d.entities))}
self.relation_idxs = {
d.relations[i]: i
for i in range(len(d.relations))
}
f = open('../data/' + self.dataset + '/entities.dict', 'w')
for key, value in self.entity_idxs.items():
f.write(key + '\t' + str(value) + '\n')
f.close()
f = open('../data/' + self.dataset + '/relations.dict', 'w')
for key, value in self.relation_idxs.items():
f.write(key + '\t' + str(value) + '\n')
f.close()
train_data_idxs = self.get_data_idxs(d.train_data)
print("Number of training data points: %d" % len(train_data_idxs))
print('Entities: %d' % len(self.entity_idxs))
print('Relations: %d' % len(self.relation_idxs))
model = TuckER(d, self.ent_vec_dim, self.rel_vec_dim, **self.kwargs)
model.init()
if self.load_from != '':
fname = self.load_from
checkpoint = torch.load(fname)
model.load_state_dict(checkpoint)
if self.cuda:
model.cuda()
opt = torch.optim.Adam(model.parameters(), lr=self.learning_rate)
if self.decay_rate:
scheduler = ExponentialLR(opt, self.decay_rate)
er_vocab = self.get_er_vocab(train_data_idxs)
er_vocab_pairs = list(er_vocab.keys())
print("Starting training...")
for it in range(1, self.num_iterations + 1):
start_train = time.time()
model.train()
losses = []
np.random.shuffle(er_vocab_pairs)
for j in tqdm(range(0, len(er_vocab_pairs), self.batch_size)):
data_batch, targets = self.get_batch(er_vocab, er_vocab_pairs,
j)
opt.zero_grad()
e1_idx = torch.tensor(data_batch[:, 0])
r_idx = torch.tensor(data_batch[:, 1])
if self.cuda:
e1_idx = e1_idx.cuda()
r_idx = r_idx.cuda()
predictions = model.forward(e1_idx, r_idx)
if self.label_smoothing:
targets = ((1.0 - self.label_smoothing) *
targets) + (1.0 / targets.size(1))
loss = model.loss(predictions, targets)
loss.backward()
opt.step()
losses.append(loss.item())
if self.decay_rate:
scheduler.step()
if it % 100 == 0:
print('Epoch', it, ' Epoch time',
time.time() - start_train, ' Loss:', np.mean(losses))
model.eval()
with torch.no_grad():
if it % self.valid_steps == 0:
start_test = time.time()
print("Validation:")
valid = self.evaluate(model, d.valid_data)
print("Test:")
test = self.evaluate(model, d.test_data)
valid_mrr = valid[0]
test_mrr = test[0]
if valid_mrr >= best_valid[0]:
best_valid = valid
best_test = test
print('Validation MRR increased.')
print('Saving model...')
write_embedding_files(model)
print('Model saved!')
print('Best valid:', best_valid)
print('Best Test:', best_test)
print('Dataset:', self.dataset)
print('Model:', self.model)
print(time.time() - start_test)
print(
'Learning rate %f | Decay %f | Dim %d | Input drop %f | Hidden drop 2 %f | LS %f | Batch size %d | Loss type %s | L3 reg %f'
%
(self.learning_rate, self.decay_rate, self.ent_vec_dim,
self.kwargs["input_dropout"],
self.kwargs["hidden_dropout2"], self.label_smoothing,
self.batch_size, self.loss_type, self.l3_reg))
def main(rank, args):
# Distributed setup
if args.distributed:
setup_distributed(rank, args.world_size)
not_main_rank = args.distributed and rank != 0
logging.info("Start time: %s", datetime.now())
# Explicitly set seed to make sure models created in separate processes
# start from same random weights and biases
torch.manual_seed(args.seed)
# Empty CUDA cache
torch.cuda.empty_cache()
# Change backend for flac files
torchaudio.set_audio_backend("soundfile")
# Transforms
melkwargs = {
"n_fft": args.win_length,
"n_mels": args.n_bins,
"hop_length": args.hop_length,
}
sample_rate_original = 16000
if args.type == "mfcc":
transforms = torch.nn.Sequential(
torchaudio.transforms.MFCC(
sample_rate=sample_rate_original,
n_mfcc=args.n_bins,
melkwargs=melkwargs,
), )
num_features = args.n_bins
elif args.type == "waveform":
transforms = torch.nn.Sequential(UnsqueezeFirst())
num_features = 1
else:
raise ValueError("Model type not supported")
if args.normalize:
transforms = torch.nn.Sequential(transforms, Normalize())
augmentations = torch.nn.Sequential()
if args.freq_mask:
augmentations = torch.nn.Sequential(
augmentations,
torchaudio.transforms.FrequencyMasking(
freq_mask_param=args.freq_mask),
)
if args.time_mask:
augmentations = torch.nn.Sequential(
augmentations,
torchaudio.transforms.TimeMasking(time_mask_param=args.time_mask),
)
# Text preprocessing
char_blank = "*"
char_space = " "
char_apostrophe = "'"
labels = char_blank + char_space + char_apostrophe + string.ascii_lowercase
language_model = LanguageModel(labels, char_blank, char_space)
# Dataset
training, validation = split_process_librispeech(
[args.dataset_train, args.dataset_valid],
[transforms, transforms],
language_model,
root=args.dataset_root,
folder_in_archive=args.dataset_folder_in_archive,
)
# Decoder
if args.decoder == "greedy":
decoder = GreedyDecoder()
else:
raise ValueError("Selected decoder not supported")
# Model
model = Wav2Letter(
num_classes=language_model.length,
input_type=args.type,
num_features=num_features,
)
if args.jit:
model = torch.jit.script(model)
if args.distributed:
n = torch.cuda.device_count() // args.world_size
devices = list(range(rank * n, (rank + 1) * n))
model = model.to(devices[0])
model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=devices)
else:
devices = ["cuda" if torch.cuda.is_available() else "cpu"]
model = model.to(devices[0], non_blocking=True)
model = torch.nn.DataParallel(model)
n = count_parameters(model)
logging.info("Number of parameters: %s", n)
# Optimizer
if args.optimizer == "adadelta":
optimizer = Adadelta(
model.parameters(),
lr=args.learning_rate,
weight_decay=args.weight_decay,
eps=args.eps,
rho=args.rho,
)
elif args.optimizer == "sgd":
optimizer = SGD(
model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay,
)
elif args.optimizer == "adam":
optimizer = Adam(
model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay,
)
elif args.optimizer == "adamw":
optimizer = AdamW(
model.parameters(),
lr=args.learning_rate,
momentum=args.momentum,
weight_decay=args.weight_decay,
)
else:
raise ValueError("Selected optimizer not supported")
if args.scheduler == "exponential":
scheduler = ExponentialLR(optimizer, gamma=args.gamma)
elif args.scheduler == "reduceonplateau":
scheduler = ReduceLROnPlateau(optimizer, patience=10, threshold=1e-3)
else:
raise ValueError("Selected scheduler not supported")
criterion = torch.nn.CTCLoss(blank=language_model.mapping[char_blank],
zero_infinity=False)
# Data Loader
collate_fn_train = collate_factory(model_length_function, augmentations)
collate_fn_valid = collate_factory(model_length_function)
loader_training_params = {
"num_workers": args.workers,
"pin_memory": True,
"shuffle": True,
"drop_last": True,
}
loader_validation_params = loader_training_params.copy()
loader_validation_params["shuffle"] = False
loader_training = DataLoader(
training,
batch_size=args.batch_size,
collate_fn=collate_fn_train,
**loader_training_params,
)
loader_validation = DataLoader(
validation,
batch_size=args.batch_size,
collate_fn=collate_fn_valid,
**loader_validation_params,
)
# Setup checkpoint
best_loss = 1.0
load_checkpoint = args.checkpoint and os.path.isfile(args.checkpoint)
if args.distributed:
torch.distributed.barrier()
if load_checkpoint:
logging.info("Checkpoint: loading %s", args.checkpoint)
checkpoint = torch.load(args.checkpoint)
args.start_epoch = checkpoint["epoch"]
best_loss = checkpoint["best_loss"]
model.load_state_dict(checkpoint["state_dict"])
optimizer.load_state_dict(checkpoint["optimizer"])
scheduler.load_state_dict(checkpoint["scheduler"])
logging.info("Checkpoint: loaded '%s' at epoch %s", args.checkpoint,
checkpoint["epoch"])
else:
logging.info("Checkpoint: not found")
save_checkpoint(
{
"epoch": args.start_epoch,
"state_dict": model.state_dict(),
"best_loss": best_loss,
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
},
False,
args.checkpoint,
not_main_rank,
)
if args.distributed:
torch.distributed.barrier()
torch.autograd.set_detect_anomaly(False)
for epoch in range(args.start_epoch, args.epochs):
logging.info("Epoch: %s", epoch)
train_one_epoch(
model,
criterion,
optimizer,
scheduler,
loader_training,
decoder,
language_model,
devices[0],
epoch,
args.clip_grad,
not_main_rank,
not args.reduce_lr_valid,
)
loss = evaluate(
model,
criterion,
loader_validation,
decoder,
language_model,
devices[0],
epoch,
not_main_rank,
)
if args.reduce_lr_valid and isinstance(scheduler, ReduceLROnPlateau):
scheduler.step(loss)
is_best = loss < best_loss
best_loss = min(loss, best_loss)
save_checkpoint(
{
"epoch": epoch + 1,
"state_dict": model.state_dict(),
"best_loss": best_loss,
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict(),
},
is_best,
args.checkpoint,
not_main_rank,
)
logging.info("End time: %s", datetime.now())
if args.distributed:
torch.distributed.destroy_process_group()
class TrackerSiamFC(Tracker):
def __init__(self, net_path=None, **kargs):
super(TrackerSiamFC, self).__init__(name='SiamFC',
is_deterministic=True)
self.cfg = self.parse_args(**kargs)
# setup GPU device if available
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
# setup model
self.net = SiamFC()
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
self.net = self.net.to(self.device)
# setup optimizer
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.cfg.initial_lr,
weight_decay=self.cfg.weight_decay,
momentum=self.cfg.momentum)
# setup lr scheduler
self.lr_scheduler = ExponentialLR(self.optimizer,
gamma=self.cfg.lr_decay)
def parse_args(self, **kargs):
# default parameters
cfg = {
# inference parameters
'exemplar_sz': 127,
'instance_sz': 255,
'context': 0.5,
'scale_num': 3,
'scale_step': 1.0375,
'scale_lr': 0.59,
'scale_penalty': 0.9745,
'window_influence': 0.176,
'response_sz': 17,
'response_up': 16,
'total_stride': 8,
'adjust_scale': 0.001,
# train parameters
'initial_lr': 0.01,
'lr_decay': 0.8685113737513527,
'weight_decay': 5e-4,
'momentum': 0.9,
'r_pos': 16,
'r_neg': 0
}
for key, val in kargs.items():
if key in cfg:
cfg.update({key: val})
return namedtuple('GenericDict', cfg.keys())(**cfg)
def init(self, image, box):
image = np.asarray(image)
# convert box to 0-indexed and center based [cy, cx, h, w]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
self.center, self.target_sz = box[:2], box[2:]
# create hanning window
self.upscale_sz = self.cfg.response_up * self.cfg.response_sz
self.hann_window = np.outer(np.hanning(self.upscale_sz),
np.hanning(self.upscale_sz))
self.hann_window /= self.hann_window.sum()
# search scale factors
self.scale_factors = self.cfg.scale_step**np.linspace(
-(self.cfg.scale_num // 2), self.cfg.scale_num // 2,
self.cfg.scale_num)
# exemplar and search sizes
context = self.cfg.context * np.sum(self.target_sz)
self.z_sz = np.sqrt(np.prod(self.target_sz + context))
self.x_sz = self.z_sz * \
self.cfg.instance_sz / self.cfg.exemplar_sz
# exemplar image
self.avg_color = np.mean(image, axis=(0, 1))
exemplar_image = self._crop_and_resize(image,
self.center,
self.z_sz,
out_size=self.cfg.exemplar_sz,
pad_color=self.avg_color)
gt_img = image[int(box[0] - box[2] / 2):int(box[0] + box[2] / 2),
int(box[1] - box[3] / 2):int(box[1] + box[3] / 2)]
self.gt_img = gt_img #added by gtz
# exemplar features
exemplar_image = torch.from_numpy(exemplar_image).to(
self.device).permute([2, 0, 1]).unsqueeze(0).float()
with torch.set_grad_enabled(False):
self.net.eval()
self.kernel = self.net.feature(exemplar_image)
self.init_box = box #added by gtz
def update(self, image):
image = np.asarray(image)
# search images
instance_images = [
self._crop_and_resize(image,
self.center,
self.x_sz * f,
out_size=self.cfg.instance_sz,
pad_color=self.avg_color)
for f in self.scale_factors
]
instance_images = np.stack(instance_images, axis=0)
instance_images = torch.from_numpy(instance_images).to(
self.device).permute([0, 3, 1, 2]).float()
# responses
with torch.set_grad_enabled(False):
self.net.eval()
instances = self.net.feature(instance_images)
responses = F.conv2d(instances, self.kernel) * 0.001
responses = responses.squeeze(1).cpu().numpy()
# upsample responses and penalize scale changes
responses = np.stack([
cv2.resize(t, (self.upscale_sz, self.upscale_sz),
interpolation=cv2.INTER_CUBIC) for t in responses
],
axis=0)
responses[:self.cfg.scale_num // 2] *= self.cfg.scale_penalty
responses[self.cfg.scale_num // 2 + 1:] *= self.cfg.scale_penalty
# peak scale
scale_id = np.argmax(np.amax(responses, axis=(1, 2)))
# peak location
response = responses[scale_id]
response -= response.min()
response /= response.sum() + 1e-16
response = (1 - self.cfg.window_influence) * response + \
self.cfg.window_influence * self.hann_window
loc = np.unravel_index(response.argmax(),
response.shape) #tuple (138,136)
# locate target center
disp_in_response = np.array(loc) - self.upscale_sz // 2
disp_in_instance = disp_in_response * \
self.cfg.total_stride / self.cfg.response_up
disp_in_image = disp_in_instance * self.x_sz * \
self.scale_factors[scale_id] / self.cfg.instance_sz
self.center += disp_in_image
# update target size
scale = (1 - self.cfg.scale_lr) * 1.0 + \
self.cfg.scale_lr * self.scale_factors[scale_id]
self.target_sz *= scale
self.z_sz *= scale
self.x_sz *= scale
# return 1-indexed and left-top based bounding box
box = np.array([
self.center[1] + 1 - (self.target_sz[1] - 1) / 2,
self.center[0] + 1 - (self.target_sz[0] - 1) / 2,
self.target_sz[1], self.target_sz[0]
])
return box
def step(self, batch, backward=True, update_lr=False):
if backward:
self.net.train()
if update_lr:
self.lr_scheduler.step()
else:
self.net.eval()
z = batch[0].to(self.device)
x = batch[1].to(self.device)
with torch.set_grad_enabled(backward):
responses = self.net(z, x)
labels, weights = self._create_labels(responses.size())
loss = F.binary_cross_entropy_with_logits(responses,
labels,
weight=weights,
size_average=True)
if backward:
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
def _crop_and_resize(self, image, center, size, out_size, pad_color):
# convert box to corners (0-indexed)
size = round(size)
corners = np.concatenate((np.round(center - (size - 1) / 2),
np.round(center - (size - 1) / 2) + size))
corners = np.round(corners).astype(int)
# pad image if necessary
pads = np.concatenate((-corners[:2], corners[2:] - image.shape[:2]))
npad = max(0, int(pads.max()))
if npad > 0:
image = cv2.copyMakeBorder(image,
npad,
npad,
npad,
npad,
cv2.BORDER_CONSTANT,
value=pad_color)
# crop image patch
corners = (corners + npad).astype(int)
patch = image[corners[0]:corners[2], corners[1]:corners[3]]
# resize to out_size
patch = cv2.resize(patch, (out_size, out_size))
return patch
def _create_labels(self, size):
# skip if same sized labels already created
if hasattr(self, 'labels') and self.labels.size() == size:
return self.labels, self.weights
def logistic_labels(x, y, r_pos, r_neg):
dist = np.abs(x) + np.abs(y) # block distance
labels = np.where(
dist <= r_pos, np.ones_like(x),
np.where(dist < r_neg,
np.ones_like(x) * 0.5, np.zeros_like(x)))
return labels
# distances along x- and y-axis
n, c, h, w = size
x = np.arange(w) - w // 2
y = np.arange(h) - h // 2
x, y = np.meshgrid(x, y)
# create logistic labels
r_pos = self.cfg.r_pos / self.cfg.total_stride
r_neg = self.cfg.r_neg / self.cfg.total_stride
labels = logistic_labels(x, y, r_pos, r_neg)
# pos/neg weights
pos_num = np.sum(labels == 1)
neg_num = np.sum(labels == 0)
weights = np.zeros_like(labels)
weights[labels == 1] = 0.5 / pos_num
weights[labels == 0] = 0.5 / neg_num
weights *= pos_num + neg_num
# repeat to size
labels = labels.reshape((1, 1, h, w))
weights = weights.reshape((1, 1, h, w))
labels = np.tile(labels, (n, c, 1, 1))
weights = np.tile(weights, [n, c, 1, 1])
# convert to tensors
self.labels = torch.from_numpy(labels).to(self.device).float()
self.weights = torch.from_numpy(weights).to(self.device).float()
return self.labels, self.weights
def visualize(self, img_files, box):
frame_num = len(img_files)
boxes = np.zeros((frame_num, 4))
boxes[0] = box
times = np.zeros(frame_num)
for f, img_file in enumerate(img_files):
image = Image.open(img_file)
if not image.mode == 'RGB':
image = image.convert('RGB')
start_time = time.time()
if f == 0:
self.init(image, box)
else:
boxes[f, :] = self.update(image) # x,y,w,h
heatmap, response = self.update_heatmap(image)
colormap, Prob1 = self.update_colormap(image)
mixedmap = self.update_mixedmap(image.size, response, Prob1)
image = np.asarray(image)
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
x1, y1, w, h = boxes[f, :]
x1 = int(x1)
y1 = int(y1)
x2 = int(x1 + w)
y2 = int(y1 + h)
cv2.rectangle(image, (x1, y1), (x2, y2), (0, 0, 255), 2)
times[f] = time.time() - start_time
return boxes, times
def visualize1(self, img_files, gt_boxes, opt, model):
for f, img_file in enumerate(img_files):
image = Image.open(img_file)
if not image.mode == 'RGB':
image = image.convert('RGB')
box = gt_boxes[f]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
center, target_sz = box[:2], box[2:]
search_size = 2 * np.max(target_sz)
start_time = time.time()
patch = self.my_crop(image, center, search_size)
patch = np.asarray(patch)
# image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
candidates = detect(opt, model, patch)[0]
self.save_imgs(opt, candidates, patch, target_sz, f)
def save_imgs(self, opt, candidates, image, target_sz, f):
# Bounding-box colors
color = (0.6, 0.3, 0.3, 1)
print("\nSaving images:")
# Create plot
plt.figure()
fig, ax = plt.subplots(1)
ax.imshow(image)
# Draw bounding boxes and labels of detections
if candidates is not None:
# Rescale boxes to original image
candidates = rescale_boxes(candidates, opt.img_size,
image.shape[:2])
candidate_sz = candidates[:, 2:4] - candidates[:, :2]
target_sz1 = torch.Tensor([target_sz[1],
target_sz[0]]).expand_as(candidate_sz)
# double filtering
valid_mask_1 = ~(torch.any(
candidate_sz > opt.one_scale_thres * target_sz1, dim=1)
| torch.any(candidate_sz <
1 / opt.one_scale_thres * target_sz1,
dim=1))
valid_mask_2 = ~(torch.all(
candidate_sz > opt.two_scale_thres * target_sz1, dim=1)
| torch.all(candidate_sz <
1 / opt.two_scale_thres * target_sz1,
dim=1))
valid_mask = valid_mask_1 & valid_mask_2
candidates = candidates[valid_mask]
for x1, y1, x2, y2, conf in candidates:
box_w = x2 - x1
box_h = y2 - y1
# Create a Rectangle patch
bbox = patches.Rectangle((x1, y1),
box_w,
box_h,
linewidth=2,
edgecolor=color,
facecolor="none")
# Add the bbox to the plot
ax.add_patch(bbox)
# Save generated image with detections
plt.axis("off")
plt.gca().xaxis.set_major_locator(NullLocator())
plt.gca().yaxis.set_major_locator(NullLocator())
plt.savefig(f"my_output1/{f}.png", bbox_inches="tight", pad_inches=0.0)
plt.close()
def my_crop(self, image, center, size):
image = np.array(image)
# convert box to corners (0-indexed)
size = round(size)
corners = np.concatenate((np.round(center - (size - 1) / 2),
np.round(center - (size - 1) / 2) + size))
corners = np.round(corners).astype(int)
# pad image if necessary
pads = np.concatenate((-corners[:2], corners[2:] - image.shape[:2]))
npad = max(0, int(pads.max()))
if npad > 0:
image = cv2.copyMakeBorder(image,
npad,
npad,
npad,
npad,
cv2.BORDER_CONSTANT,
value=(255, 255, 255))
# crop image patch
corners = (corners + npad).astype(int)
patch = image[corners[0]:corners[2], corners[1]:corners[3]]
return patch
def update_colormap(self, image):
image = np.asarray(image)
hist_obj = cv2.calcHist([self.gt_img], [0, 1, 2], None, [10, 10, 10],
[0, 256, 0, 256, 0, 256])
# cv2.imshow('gt_img',self.gt_img)
# cv2.waitKey(0)
hist_img = cv2.calcHist([image], [0, 1, 2], None, [10, 10, 10],
[0, 256, 0, 256, 0, 256])
image_10 = image / 25.6 # histogram has 10 bins
image_10 = image_10.astype('uint8')
# creating a likelihood image acc. to obj-surr or obj-distractor model
a = image_10[:, :, 0]
a = a.ravel()
b = image_10[:, :, 1]
b = b.ravel()
c_ = image_10[:, :, 2]
c_ = c_.ravel()
H_obj = hist_obj[
a, b,
c_] # image with pixel value=bin count of the pixel value at the same location in original image
H_img = hist_img[a, b, c_]
Prob1 = np.zeros((image.shape[0] * image.shape[1], ), dtype='float')
H_obj = H_obj.astype('float')
H_img = H_img.astype('float')
mask = H_img == 0
# print mask,"check itjhjnkjkjkjk"
Prob1[~mask] = np.divide(H_obj[~mask], H_img[~mask])
Prob1[mask] = 0.1
Prob1 = Prob1.reshape((image.shape[0], image.shape[1]))
Prob2 = (Prob1) * 255
Prob2 = Prob2.astype('uint8')
likemap = cv2.applyColorMap(Prob2, cv2.COLORMAP_JET)
return likemap, Prob1
def update_mixedmap(self, shape, response, Prob1):
#response is already normalized
# Prob1-=Prob1.min() # Prob1 -> (432,576)
# Prob1/=Prob1.sum()+1e-16
#response = np.expand_dims(response, -1).astype('uint8') # (272,272) -> (272,272,1)
response = cv2.resize(response, (shape[0], shape[1]),
interpolation=cv2.INTER_CUBIC)
enhance_mask = Prob1 > 0.5
slight_reduce_mask = Prob1 < 1e-2
remove_mask = Prob1 < 1e-4
#mixedmap=response-reduce_mask*0.5e-5
#mixedmap=response+enhance_mask*1e-6-slight_reduce_mask*8e-6-remove_mask*2e-5
mixedmap = response - slight_reduce_mask * 1e-5 - remove_mask * 5e-5
mixedmap = np.clip(mixedmap, 0, 1)
m_max = np.max(mixedmap)
m_min = np.min(mixedmap)
mixedmap = (mixedmap - m_min) / (m_max - m_min) * 255.
mixedmap = mixedmap.astype('uint8')
mixedmap = cv2.applyColorMap(mixedmap, cv2.COLORMAP_JET)
return mixedmap
class TrackerSiamFC(Tracker):
def __init__(self, net_path=None, **kargs):
super(TrackerSiamFC, self).__init__(name='SiamFC',
is_deterministic=True)
self.cfg = self.parse_args(**kargs)
# setup GPU device if available
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
# setup model
self.net = SiamFC()
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
self.net = self.net.to(self.device)
# setup optimizer
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.cfg.initial_lr,
weight_decay=self.cfg.weight_decay,
momentum=self.cfg.momentum)
# setup lr scheduler
self.lr_scheduler = ExponentialLR(self.optimizer,
gamma=self.cfg.lr_decay)
def parse_args(self, **kargs):
# default parameters
cfg = {
# inference parameters
'exemplar_sz': 135,
'instance_sz': 263,
'context': 0.5,
'scale_num': 3,
'scale_step': 1.0375,
'scale_lr': 0.59,
'scale_penalty': 0.9745,
'window_influence': 0.27, #0.176,
'response_sz': 17,
'response_up': 16,
'total_stride': 8,
'adjust_scale': 0.001,
# train parameters
'initial_lr': 0.01,
'lr_decay': 0.8685113737513527,
'weight_decay': 5e-4,
'momentum': 0.9,
'r_pos': 16,
'r_neg': 0
}
for key, val in kargs.items():
if key in cfg:
cfg.update({key: val})
return namedtuple('GenericDict', cfg.keys())(**cfg)
def init(self, image, box):
image = np.asarray(image)
# convert box to 0-indexed and center based [y, x, h, w]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
self.center, self.target_sz = box[:2], box[2:]
# create hanning window
self.upscale_sz = self.cfg.response_up * self.cfg.response_sz
self.hann_window = np.outer(np.hanning(self.upscale_sz),
np.hanning(self.upscale_sz))
self.hann_window /= self.hann_window.sum()
# search scale factors
self.scale_factors = self.cfg.scale_step**np.linspace(
-(self.cfg.scale_num // 2), self.cfg.scale_num // 2,
self.cfg.scale_num)
# exemplar and search sizes
context = self.cfg.context * np.sum(self.target_sz)
self.z_sz = np.sqrt(np.prod(self.target_sz + context))
self.x_sz = self.z_sz * \
self.cfg.instance_sz / self.cfg.exemplar_sz
# exemplar image
self.avg_color = np.mean(image, axis=(0, 1))
exemplar_image = self._crop_and_resize(image,
self.center,
self.z_sz,
out_size=self.cfg.exemplar_sz,
pad_color=self.avg_color)
# exemplar features
exemplar_image = torch.from_numpy(exemplar_image).to(
self.device).permute([2, 0, 1]).unsqueeze(0).float()
with torch.set_grad_enabled(False):
self.net.eval()
z = self.net.feature1(exemplar_image)
z_noise = self.net.feature2(exemplar_image)
z = torch.add(z, z_noise)
self.kernel = self.net.feature3(z)
def update(self, image):
image = np.asarray(image)
# search images
instance_images = [
self._crop_and_resize(image,
self.center,
self.x_sz * f,
out_size=self.cfg.instance_sz,
pad_color=self.avg_color)
for f in self.scale_factors
]
instance_images = np.stack(instance_images, axis=0)
instance_images = torch.from_numpy(instance_images).to(
self.device).permute([0, 3, 1, 2]).float()
# responses
with torch.set_grad_enabled(False):
self.net.eval()
x = self.net.feature1(instance_images)
x_noise = self.net.feature2(instance_images)
x = torch.add(x, x_noise)
instances = self.net.feature3(x)
responses = F.conv2d(instances, self.kernel) * 0.001
responses = responses.squeeze(1).cpu().numpy()
# upsample responses and penalize scale changes
responses = np.stack([
cv2.resize(t, (self.upscale_sz, self.upscale_sz),
interpolation=cv2.INTER_CUBIC) for t in responses
],
axis=0)
responses[:self.cfg.scale_num // 2] *= self.cfg.scale_penalty
responses[self.cfg.scale_num // 2 + 1:] *= self.cfg.scale_penalty
# peak scale
scale_id = np.argmax(np.amax(responses, axis=(1, 2)))
# peak location
response = responses[scale_id]
response -= response.min()
response /= response.sum() + 1e-16
response = (1 - self.cfg.window_influence) * response + \
self.cfg.window_influence * self.hann_window
loc = np.unravel_index(response.argmax(), response.shape)
# locate target center
disp_in_response = np.array(loc) - self.upscale_sz // 2
disp_in_instance = disp_in_response * \
self.cfg.total_stride / self.cfg.response_up
disp_in_image = disp_in_instance * self.x_sz * \
self.scale_factors[scale_id] / self.cfg.instance_sz
self.center += disp_in_image
# update target size
scale = (1 - self.cfg.scale_lr) * 1.0 + \
self.cfg.scale_lr * self.scale_factors[scale_id]
self.target_sz *= scale
self.z_sz *= scale
self.x_sz *= scale
# return 1-indexed and left-top based bounding box
box = np.array([
self.center[1] + 1 - (self.target_sz[1] - 1) / 2,
self.center[0] + 1 - (self.target_sz[0] - 1) / 2,
self.target_sz[1], self.target_sz[0]
])
return box
def step(self, batch, backward=True, update_lr=False):
if backward:
self.net.train()
if update_lr:
self.lr_scheduler.step()
else:
self.net.eval()
z = batch[0].to(self.device)
z_noise = batch[1].to(self.device)
x = batch[2].to(self.device)
x_noise = batch[3].to(self.device)
with torch.set_grad_enabled(backward):
responses = self.net(z, z_noise, x, x_noise)
labels, weights = self._create_labels(responses.size())
loss = F.binary_cross_entropy_with_logits(responses,
labels,
weight=weights,
size_average=True)
if backward:
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
#print(loss.item())
return loss.item()
def _crop_and_resize(self, image, center, size, out_size, pad_color):
# convert box to corners (0-indexed)
size = round(size)
corners = np.concatenate((np.round(center - (size - 1) / 2),
np.round(center - (size - 1) / 2) + size))
corners = np.round(corners).astype(int)
# pad image if necessary
pads = np.concatenate((-corners[:2], corners[2:] - image.shape[:2]))
npad = max(0, int(pads.max()))
if npad > 0:
image = cv2.copyMakeBorder(image,
npad,
npad,
npad,
npad,
cv2.BORDER_CONSTANT,
value=pad_color)
# crop image patch
corners = (corners + npad).astype(int)
patch = image[corners[0]:corners[2], corners[1]:corners[3]]
# resize to out_size
patch = cv2.resize(patch, (out_size, out_size))
return patch
def _create_labels(self, size):
# skip if same sized labels already created
if hasattr(self, 'labels') and self.labels.size() == size:
return self.labels, self.weights
def logistic_labels(x, y, r_pos, r_neg):
dist = np.abs(x) + np.abs(y) # block distance
labels = np.where(
dist <= r_pos, np.ones_like(x),
np.where(dist < r_neg,
np.ones_like(x) * 0.5, np.zeros_like(x)))
return labels
# distances along x- and y-axis
n, c, h, w = size
x = np.arange(w) - w // 2
y = np.arange(h) - h // 2
x, y = np.meshgrid(x, y)
# create logistic labels
r_pos = self.cfg.r_pos / self.cfg.total_stride
r_neg = self.cfg.r_neg / self.cfg.total_stride
labels = logistic_labels(x, y, r_pos, r_neg)
# pos/neg weights
pos_num = np.sum(labels == 1)
neg_num = np.sum(labels == 0)
weights = np.zeros_like(labels)
weights[labels == 1] = 0.5 / pos_num
weights[labels == 0] = 0.5 / neg_num
weights *= pos_num + neg_num
# repeat to size
labels = labels.reshape((1, 1, h, w))
weights = weights.reshape((1, 1, h, w))
labels = np.tile(labels, (n, c, 1, 1))
weights = np.tile(weights, [n, c, 1, 1])
# convert to tensors
self.labels = torch.from_numpy(labels).to(self.device).float()
self.weights = torch.from_numpy(weights).to(self.device).float()
return self.labels, self.weights
def run(train_loader, val_loader, epochs, lr, momentum, weight_decay, lr_step,
k1, k2, es_patience, log_dir):
model = Vgg16()
device = 'cpu'
if torch.cuda.is_available():
device = 'cuda'
model.to(device)
optimizer = optim.SGD(model.parameters(),
lr=lr,
momentum=momentum,
weight_decay=weight_decay)
lr_scheduler = ExponentialLR(optimizer, gamma=0.975)
# criterion = VAELoss(k1=k1, k2=k2).to(device)
def update_fn(engine, batch):
x, y = _prepare_batch(batch, device=device, non_blocking=True)
model.train()
optimizer.zero_grad()
output = model(x)
# Compute loss
loss = F.nll_loss(output, y)
loss.backward()
optimizer.step()
return {
"batchloss": loss.item(),
}
trainer = Engine(update_fn)
try:
GpuInfo().attach(trainer)
except RuntimeError:
print(
"INFO: By default, in this example it is possible to log GPU information (used memory, utilization). "
"As there is no pynvml python package installed, GPU information won't be logged. Otherwise, please "
"install it : `pip install pynvml`")
trainer.add_event_handler(Events.ITERATION_COMPLETED(every=lr_step),
lambda engine: lr_scheduler.step())
metric_names = [
'batchloss',
]
def output_transform(x, name):
return x[name]
for n in metric_names:
# We compute running average values on the output (batch loss) across all devices
RunningAverage(output_transform=partial(output_transform, name=n),
epoch_bound=False,
device=device).attach(trainer, n)
exp_name = datetime.now().strftime("%Y%m%d-%H%M%S")
log_path = log_dir + "/vgg_vae/{}".format(exp_name)
tb_logger = TensorboardLogger(log_dir=log_path)
tb_logger.attach(trainer,
log_handler=OutputHandler(tag="training",
metric_names=metric_names),
event_name=Events.ITERATION_COMPLETED)
tb_logger.attach(trainer,
log_handler=OptimizerParamsHandler(optimizer, "lr"),
event_name=Events.ITERATION_STARTED)
ProgressBar(persist=True,
bar_format="").attach(trainer,
event_name=Events.EPOCH_STARTED,
closing_event_name=Events.COMPLETED)
ProgressBar(persist=False, bar_format="").attach(trainer,
metric_names=metric_names)
# val process definition
def loss_output_transform(output):
return output
def acc_output_transform(output):
return output
customed_loss = Loss(loss_fn=F.nll_loss,
output_transform=loss_output_transform,
device=device)
customed_accuracy = Accuracy(output_transform=acc_output_transform,
device=device)
metrics = {'Loss': customed_loss, 'Accuracy': customed_accuracy}
def val_update_fn(engine, batch):
model.eval()
with torch.no_grad():
x, y = _prepare_batch(batch, device=device, non_blocking=True)
output = model(x)
return output, y
val_evaluator = Engine(val_update_fn)
for name, metric in metrics.items():
metric.attach(val_evaluator, name)
def run_evaluation(engine):
val_evaluator.run(val_loader)
trainer.add_event_handler(Events.EPOCH_COMPLETED, run_evaluation)
trainer.add_event_handler(Events.COMPLETED, run_evaluation)
ProgressBar(persist=False, desc="Train evaluation").attach(val_evaluator)
# Log val metrics:
tb_logger.attach(val_evaluator,
log_handler=OutputHandler(tag="val",
metric_names=list(
metrics.keys()),
another_engine=trainer),
event_name=Events.EPOCH_COMPLETED)
# trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())
# Store the best model
def default_score_fn(engine):
score = engine.state.metrics['Accuracy']
return score
best_model_handler = ModelCheckpoint(dirname=log_path,
filename_prefix="best",
n_saved=3,
score_name="val_acc",
score_function=default_score_fn)
val_evaluator.add_event_handler(Events.COMPLETED, best_model_handler, {
'model': model,
})
# Add early stopping
es_patience = es_patience
es_handler = EarlyStopping(patience=es_patience,
score_function=default_score_fn,
trainer=trainer)
val_evaluator.add_event_handler(Events.COMPLETED, es_handler)
setup_logger(es_handler._logger)
setup_logger(logging.getLogger("ignite.engine.engine.Engine"))
def empty_cuda_cache(engine):
torch.cuda.empty_cache()
import gc
gc.collect()
trainer.add_event_handler(Events.EPOCH_COMPLETED, empty_cuda_cache)
val_evaluator.add_event_handler(Events.COMPLETED, empty_cuda_cache)
trainer.run(train_loader, max_epochs=epochs)
# optimizer.zero_grad()
# loss.backward()
# optimizer.step()
#
# losses3 += loss.item()
if local_rank == 0:
total = i + 1
tbar.set_description(
'epoch: %d, loss manual mining: %.3f, loss hard mining: %.3f'
% (epoch + 1, losses_manual_mining / total,
losses_hard_mining / total))
# tbar.set_description('epoch: %d, loss1: %.3f, loss2: %.3f'
# % (epoch + 1, losses_manual_mining / (i + 1), losses_hard_mining / (i + 1)))
scheduler.step(epoch)
if local_rank == 0:
checkpoints = {
'query': text_encoder.module.state_dict(),
'item': image_encoder.module.state_dict(),
'score': score_model.module.state_dict(),
# 'category': category_embedding.state_dict(),
# 'generator': text_generator.state_dict(),
'optimizer': optimizer.state_dict()
}
torch.save(
checkpoints,
os.path.join(checkpoints_dir,
'model-epoch{}.pth'.format(epoch + 1)))
# score_model.eval()
# text_encoder.eval()
def train_and_eval(self):
print("Training the %s model..." % self.model_name)
outfolder = "/afs/inf.ed.ac.uk/group/project/Knowledge_Bases/stats/"
fname = "result_%dmil_%s_ws%d_w%d_c%d_d%d_" % (int(np.ceil(d.cutoff/1e6)),
self.model_name, d.window_size, int(self.w_reg*10),
int(self.c_reg*10), self.embeddings_dim)
if self.model_name.lower() == "p2v-l":
model = P2VL(self.embeddings_dim)
elif self.model_name.lower() == "p2v-p":
model = P2VP(self.embeddings_dim)
if self.cuda:
model.cuda()
opt = torch.optim.Adam(model.parameters(), lr=self.learning_rate)
if self.decay_rate:
scheduler = ExponentialLR(opt, self.decay_rate)
data_idxs = range(len(d.word_counts))
data_w_probs = [d.w_probs[i] for i in data_idxs]
data_c_probs = [d.c_probs[i] for i in data_idxs]
cooccurrences = list(d.cooccurrence_counts.keys())
def pmi_ii(data_batch, i):
return np.maximum(0.1, np.log([d.cooccurrence_counts.get((pair[i], pair[i]), min_p)/
d.w_probs[pair[i]]**2 for pair in data_batch]))
def pmi(data_batch):
targets = [d.cooccurrence_counts.get(pair, 0.)/(d.w_probs[pair[0]]*
d.c_probs[pair[1]]) for pair in data_batch]
return np.array([np.log(target) if target!=0 else -1. for target in targets])
ww_cooccurrences = [d.cooccurrence_counts.get((idx, idx), 100.) for idx in d.w_probs]
min_p = np.min(ww_cooccurrences)/1.5
losses = []
tick = time.time()
counter = 0
for i in range(1, self.num_iterations+1):
model.train()
np.random.shuffle(cooccurrences)
num_batches = int(np.ceil(len(cooccurrences)/float(self.batch_size)))
num_neg_samples = self.batch_size*self.corrupt_size*num_batches
all_neg_pairs = list(zip(np.random.choice(data_idxs, num_neg_samples, p=data_w_probs),
np.random.choice(data_idxs, num_neg_samples, p=data_c_probs)))
epoch_loss = []
for j in range(0, len(cooccurrences), self.batch_size):
counter += 1
pos_pairs = cooccurrences[j:min(j+self.batch_size, len(cooccurrences))]
neg_pairs = all_neg_pairs[j*self.corrupt_size:j*self.corrupt_size+
self.batch_size*self.corrupt_size]
data_batch = pos_pairs + neg_pairs
targets = torch.FloatTensor(pmi(data_batch))
if self.model_name.lower() == "p2v-l":
targets_w = torch.FloatTensor(np.sqrt(pmi_ii(data_batch, 0)))
targets_c = torch.FloatTensor(np.sqrt(pmi_ii(data_batch, 1)))
elif self.model_name.lower() == "p2v-p":
targets_w = torch.FloatTensor(pmi_ii(data_batch, 0))
targets_c = torch.FloatTensor(np.zeros(len(targets)))
opt.zero_grad()
data_batch = np.array(data_batch)
w_idx = torch.tensor(data_batch[:,0])
c_idx = torch.tensor(data_batch[:,1])
if self.cuda:
w_idx = w_idx.cuda()
c_idx = c_idx.cuda()
targets = targets.cuda()
targets_w = targets_w.cuda()
targets_c = targets_c.cuda()
preds, preds_w, preds_c = model.forward(w_idx, c_idx)
loss = model.loss(preds, targets) +\
self.w_reg * model.loss(preds_w, targets_w) +\
self.c_reg * model.loss(preds_c, targets_c)
loss.backward()
opt.step()
epoch_loss.append(loss.item())
if self.decay_rate and not counter%500:
scheduler.step()
print("Iteration: %d" % i)
print("Loss: %.4f" % np.mean(epoch_loss))
if not i%10:
np_W = model.W.weight.detach().cpu().numpy()
np_C = model.C.weight.detach().cpu().numpy()
if i == 10:
np.save("%s%s%d.npy" % (outfolder, fname, i), {"W":np_W, "C":np_C,
"p_w":d.w_probs, "p_c":d.c_probs, "p_wc":d.cooccurrence_counts,
"losses":losses, "i2w":d.idx_to_word})
else:
np.save("%s%s%d.npy" % (outfolder, fname, i), {"W":np_W, "C":np_C,
"losses":losses})
losses.append(np.mean(epoch_loss))
print("Time: ", str(time.time()-tick))
def train_stand(self, train_data, valid_data, derived_struct, rela_cluster,
mrr):
self.rela_to_dict(rela_cluster)
#self.args.perf_file = os.path.join(self.args.out_dir, self.args.dataset + '_std_' + str(self.args.m) + "_" + str(self.args.n) + "_" + str(mrr) + '.txt')
#plot_config(self.args)
head, tail, rela = train_data
n_train = len(head)
if self.args.optim == 'adam' or self.args.optim == 'Adam':
self.optimizer = Adam(self.model.parameters(), lr=self.args.lr)
elif self.args.optim == 'adagrad' or self.args.optim == 'Adagrad':
self.optimizer = Adagrad(self.model.parameters(), lr=self.args.lr)
else:
self.optimizer = SGD(self.model.parameters(), lr=self.args.lr)
scheduler = ExponentialLR(self.optimizer, self.args.decay_rate)
n_batch = self.args.n_batch
best_mrr = 0
start = time.time()
for epoch in range(self.args.n_stand_epoch):
#self.epoch = epoch
rand_idx = torch.randperm(n_train)
if self.GPU:
head = head[rand_idx].cuda()
tail = tail[rand_idx].cuda()
rela = rela[rand_idx].cuda()
else:
head = head[rand_idx]
tail = tail[rand_idx]
rela = rela[rand_idx]
epoch_loss = 0
n_iters = 0
#lr = scheduler.get_lr()[0]
# train model weights
for h, t, r in batch_by_size(n_batch,
head,
tail,
rela,
n_sample=n_train):
self.model.zero_grad()
loss = self.model.forward(derived_struct, h, t, r,
self.cluster_rela_dict)
loss += self.args.lamb * self.model.regul
loss.backward()
self.optimizer.step()
self.prox_operator()
epoch_loss += loss.data.cpu().numpy()
n_iters += 1
scheduler.step()
print("Epoch: %d/%d, Loss=%.2f, Stand Time=%.2f" %
(epoch + 1, self.args.n_stand_epoch, time.time() - start,
epoch_loss / n_train))
if (epoch + 1) % 5 == 0:
test, randint = True, None
valid_mrr, valid_mr, valid_1, valid_3, valid_10 = self.tester_val(
derived_struct, test, randint)
test_mrr, test_mr, test_1, test_3, test_10 = self.tester_tst(
derived_struct, test, randint)
out_str = '%d \t %.2f \t %.2f \t %.4f %.1f %.4f %.4f %.4f\t%.4f %.1f %.4f %.4f %.4f\n' % (epoch, self.time_tot, epoch_loss/n_train,\
valid_mrr, valid_mr, valid_1, valid_3, valid_10, \
test_mrr, test_mr, test_1, test_3, test_10)
# output the best performance info
if test_mrr > best_mrr:
best_mrr = test_mrr
best_str = out_str
with open(self.args.perf_file, 'a+') as f:
f.write(out_str)
with open(self.args.perf_file, 'a+') as f:
f.write("best performance:" + best_str + "\n")
f.write("struct:" + str(derived_struct) + "\n")
f.write("rela:" + str(rela_cluster) + "\n")
return best_mrr
def train_oas(self, train_data, valid_data, derived_struct):
head, tail, rela = train_data
n_train = len(head)
if self.args.optim == 'adam' or self.args.optim == 'Adam':
self.optimizer = Adam(self.model.parameters(), lr=self.args.lr)
elif self.args.optim == 'adagrad' or self.args.optim == 'Adagrad':
self.optimizer = Adagrad(self.model.parameters(), lr=self.args.lr)
else:
self.optimizer = SGD(self.model.parameters(), lr=self.args.lr)
scheduler = ExponentialLR(self.optimizer, self.args.decay_rate)
n_batch = self.args.n_batch
for epoch in range(self.args.n_oas_epoch):
start = time.time()
rand_idx = torch.randperm(n_train)
if self.GPU:
head = head[rand_idx].cuda()
tail = tail[rand_idx].cuda()
rela = rela[rand_idx].cuda()
else:
head = head[rand_idx]
tail = tail[rand_idx]
rela = rela[rand_idx]
epoch_loss = 0
n_iters = 0
# train model weights
for h, t, r in batch_by_size(n_batch,
head,
tail,
rela,
n_sample=n_train):
self.model.zero_grad()
loss = self.model.forward(derived_struct, h, t, r,
self.cluster_rela_dict)
loss += self.args.lamb * self.model.regul
loss.backward()
nn.utils.clip_grad_norm_(self.model.parameters(),
self.args.grad_clip)
self.optimizer.step()
self.prox_operator()
epoch_loss += loss.data.cpu().numpy()
n_iters += 1
scheduler.step()
if self.cluster_way == "scu":
self.rela_cluster = self.cluster()
self.rela_cluster_history.append(self.rela_cluster)
self.rela_to_dict(self.rela_cluster)
# train controller
self.train_controller()
# derive structs
self.time_tot += time.time(
) - start # evaluation for the derived architecture is unnessary in searching procedure
derived_struct, test_mrr = self.derive(sample_num=1)
print(
"Epoch: %d/%d, Search Time=%.2f, Loss=%.2f, Sampled Val MRR=%.8f, Tst MRR=%.8f"
%
(epoch + 1, self.args.n_oas_epoch, self.time_tot, epoch_loss /
n_train, self.derived_raward_history[-1], test_mrr))
class TrackerSiamFC(Tracker): #定义一个追踪器
def __init__(self, net_path=None, **kwargs):
super(TrackerSiamFC, self).__init__(net_path, True)
self.cfg = self.parse_args(**kwargs) #加载一些初始化的参数
# setup GPU device if available
self.cuda = torch.cuda.is_available()
self.device = torch.device(
'cuda:0' if self.cuda else 'cpu') #指定 GPU0 来进行训练
# setup model
self.net = Net(backbone=AlexNetV1(), head=SiamFC(self.cfg.out_scale))
ops.init_weights(self.net) #对网络权重进行初始化？？？
# load checkpoint if provided
if net_path is not None: #加载训练好的网络模型
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
self.net = self.net.to(self.device) #将模型加载到GPU上
# setup criterion
self.criterion = BalancedLoss()
# setup optimizer
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.cfg.initial_lr,
weight_decay=self.cfg.weight_decay,
momentum=self.cfg.momentum)
# setup lr scheduler #动态调整学习率 这个是怎么计算的？？
gamma = np.power(self.cfg.ultimate_lr / self.cfg.initial_lr,
1.0 / self.cfg.epoch_num)
#gamma=0.87
self.lr_scheduler = ExponentialLR(self.optimizer, gamma)
def parse_args(self, **kwargs):
# default parameters
cfg = {
# basic parameters
'out_scale': 0.001,
'exemplar_sz': 127, #第一帧
'instance_sz': 255, #
'context': 0.5,
# inference parameters
'scale_num': 3,
'scale_step': 1.0375,
'scale_lr': 0.59,
'scale_penalty': 0.9745,
'window_influence': 0.176,
'response_sz': 17,
'response_up': 16,
'total_stride': 8,
# train parameters
'epoch_num': 50,
'batch_size': 8,
'num_workers': 8, #原来是32 被我修改成8
'initial_lr': 1e-2, #0.01
'ultimate_lr': 1e-5, #0.000052
'weight_decay': 5e-4, #正则参数
'momentum': 0.9,
'r_pos': 16, #
'r_neg': 0
} #
for key, val in kwargs.items():
if key in cfg:
cfg.update({key: val})
return namedtuple('Config', cfg.keys())(**cfg)
#namedtuple比tuple更强大，与list不同的是,你不能改变tuple中元素的数值
#Namedtuple比普通tuple具有更好的可读性，可以使代码更易于维护
#为了构造一个namedtuple需要两个参数，分别是tuple的名字和其中域的名字
'''禁止计算局部梯度
方法1 使用装饰器 @torch.no_gard()修饰的函数，在调用时不允许计算梯度
方法2 # 将不用计算梯度的变量放在 with torch.no_grad()里
'''
@torch.no_grad() #
def init(self, img, box):
# set to evaluation mode
self.net.eval()
# convert box to 0-indexed and center based [y, x, h, w]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
self.center, self.target_sz = box[:2], box[
2:] #最原始的图片大小 从groundtruth读取
# create hanning window response_up=16 ； response_sz=17 ； self.upscale_sz=272
self.upscale_sz = self.cfg.response_up * self.cfg.response_sz
self.hann_window = np.outer( # np.outer 如果a，b是高维数组，函数会自动将其flatten成1维 ，用来求外积
np.hanning(self.upscale_sz), np.hanning(self.upscale_sz))
self.hann_window /= self.hann_window.sum() #？？？
# search scale factors
self.scale_factors = self.cfg.scale_step**np.linspace( #linspace 在start和stop之间返回均匀间隔的数据
-(self.cfg.scale_num // 2), #//py3中双斜杠代表向下取整
self.cfg.scale_num // 2,
self.cfg.scale_num)
# exemplar and search sizes self.cfg.context=1/2
context = self.cfg.context * np.sum(
self.target_sz) # 引入margin：2P=(长+宽）× 1/2
self.z_sz = np.sqrt(
np.prod(self.target_sz +
context)) # ([长，宽]+2P) x 2 添加 padding 没有乘以缩放因子
self.x_sz = self.z_sz * self.cfg.instance_sz / self.cfg.exemplar_sz # 226 没有乘以缩放因子
# z是初始模板的大小 x是搜索区域
# exemplar image
self.avg_color = np.mean(img,
axis=(0, 1)) # 计算RGB通道的均值,使用图像均值进行padding
z = ops.crop_and_resize(img,
self.center,
self.z_sz,
out_size=self.cfg.exemplar_sz,
border_value=self.avg_color)
#对所有的图片进行预处理，得到127x127大小的patch
# exemplar features
z = torch.from_numpy(z).to(self.device).permute(
2, 0, 1).unsqueeze(0).float()
self.kernel = self.net.backbone(z)
@torch.no_grad() #
def update(self, img):
# set to evaluation mode
self.net.eval()
# search images
x = [
ops.crop_and_resize(img,
self.center,
self.x_sz * f,
out_size=self.cfg.instance_sz,
border_value=self.avg_color)
for f in self.scale_factors
]
x = np.stack(x, axis=0)
x = torch.from_numpy(x).to(self.device).permute(0, 3, 1, 2).float()
# responses
x = self.net.backbone(x)
responses = self.net.head(self.kernel, x)
responses = responses.squeeze(1).cpu().numpy()
# upsample responses and penalize scale changes
responses = np.stack([
cv2.resize(u, (self.upscale_sz, self.upscale_sz),
interpolation=cv2.INTER_CUBIC) for u in responses
])
responses[:self.cfg.scale_num // 2] *= self.cfg.scale_penalty
responses[self.cfg.scale_num // 2 + 1:] *= self.cfg.scale_penalty
# peak scale
scale_id = np.argmax(np.amax(responses, axis=(1, 2)))
# peak location
response = responses[scale_id]
response -= response.min()
response /= response.sum() + 1e-16
response = (1 - self.cfg.window_influence) * response + \
self.cfg.window_influence * self.hann_window
loc = np.unravel_index(response.argmax(), response.shape)
# locate target center
disp_in_response = np.array(loc) - (self.upscale_sz - 1) / 2
disp_in_instance = disp_in_response * \
self.cfg.total_stride / self.cfg.response_up
disp_in_image = disp_in_instance * self.x_sz * \
self.scale_factors[scale_id] / self.cfg.instance_sz
self.center += disp_in_image
# update target size
scale = (1 - self.cfg.scale_lr
) * 1.0 + self.cfg.scale_lr * self.scale_factors[scale_id]
self.target_sz *= scale
self.z_sz *= scale
self.x_sz *= scale
# return 1-indexed and left-top based bounding box [x,y,w,h]
box = np.array([
self.center[1] + 1 - (self.target_sz[1] - 1) / 2,
self.center[0] + 1 - (self.target_sz[0] - 1) / 2,
self.target_sz[1], self.target_sz[0]
])
return box
def track(self, img_files, box, visualize=False): # x,y,w,h
frame_num = len(img_files)
boxes = np.zeros((frame_num, 4))
boxes[0] = box
times = np.zeros(frame_num)
for f, img_file in enumerate(img_files):
img = ops.read_image(img_file)
begin = time.time()
if f == 0:
self.init(img, box)
else:
boxes[f, :] = self.update(img)
times[f] = time.time() - begin
if visualize:
ops.show_image(img, boxes[f, :])
return boxes, times
def train_step(self, batch, backward=True):
# set network mode
self.net.train(backward) # 训练模式
# parse batch data
z = batch[0].to(self.device, non_blocking=self.cuda)
x = batch[1].to(self.device, non_blocking=self.cuda)
with torch.set_grad_enabled(backward):
# inference
responses = self.net(z, x)
# calculate loss
labels = self._create_labels(responses.size())
loss = self.criterion(responses, labels)
if backward:
# back propagation
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
# 在禁止计算梯度下调用被允许计算梯度的函数
@torch.enable_grad()
def train_over(self, seqs, val_seqs=None, save_dir='pretrained'):
# set to train mode
self.net.train()
# create save_dir folder
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# setup dataset
transforms = SiamFCTransforms(
exemplar_sz=self.cfg.exemplar_sz, #127
instance_sz=self.cfg.instance_sz, #255
context=self.cfg.context) # 0.5 ？？？
dataset = Pair(seqs=seqs, transforms=transforms)
# setup dataloader
dataloader = DataLoader(dataset,
batch_size=self.cfg.batch_size,
shuffle=True,
num_workers=self.cfg.num_workers,
pin_memory=self.cuda,
drop_last=True)
# loop over epochs
for epoch in range(self.cfg.epoch_num):
# update lr at each epoch
self.lr_scheduler.step(epoch=epoch)
# loop over dataloader
for it, batch in enumerate(dataloader):
loss = self.train_step(batch, backward=True)
print('Epoch: {} [{}/{}] Loss: {:.5f}'.format(
epoch + 1, it + 1, len(dataloader), loss))
sys.stdout.flush()
# save checkpoint
if not os.path.exists(save_dir):
os.makedirs(save_dir)
net_path = os.path.join(save_dir,
'siamfc_alexnet_e%d.pth' % (epoch + 1))
torch.save(self.net.state_dict(), net_path)
def _create_labels(self, size):
# skip if same sized labels already created
if hasattr(self, 'labels') and self.labels.size() == size:
return self.labels
def logistic_labels(x, y, r_pos, r_neg):
dist = np.abs(x) + np.abs(y) # block distance
labels = np.where(
dist <= r_pos, np.ones_like(x),
np.where(dist < r_neg,
np.ones_like(x) * 0.5, np.zeros_like(x)))
return labels
# distances along x- and y-axis
n, c, h, w = size
x = np.arange(w) - (w - 1) / 2
y = np.arange(h) - (h - 1) / 2
x, y = np.meshgrid(x, y)
# create logistic labels
r_pos = self.cfg.r_pos / self.cfg.total_stride
r_neg = self.cfg.r_neg / self.cfg.total_stride
labels = logistic_labels(x, y, r_pos, r_neg)
# repeat to size
labels = labels.reshape((1, 1, h, w))
labels = np.tile(labels, (n, c, 1, 1))
# convert to tensors
self.labels = torch.from_numpy(labels).to(self.device).float()
return self.labels
class Trainer_new(Trainer):
def __init__(self, args, logger=None):
self.args = args
self.logger = logger
self.best_dev_performance = 0.0
self.device = torch.device('cuda' if args.use_cuda else 'cpu')
self.learning_rate = self.args.lr
self.decay_rate = self.args.decay_rate
self.reset_time = 0
self.data_load_kg_rs()
self.load_pretrain(args)
# self.model.freeze_part()
def optim_def(self, lr):
self.learning_rate = lr
self.optim = optim.Adam(self.model.parameters(), lr=lr)
if self.decay_rate: #decay_rate > 0
self.scheduler = ExponentialLR(self.optim, self.decay_rate)
self.reset_time += 1
def model_def(self):
self.entity_total = len(self.e_map)
self.relation_total = len(self.r_map)
self.user_total = len(self.u_map)
self.item_total = len(self.i_map)
self.share_total = len(self.i_kg_map)
self.model = init_model(self.args,
self.user_total,
self.item_total,
self.entity_total,
self.relation_total,
self.logger,
None,
None,
None,
share_total=self.share_total)
def loss_def(self):
self.loss_func_kg = torch.nn.BCEWithLogitsLoss(reduction="none") #
self.loss_reg = Regularization(self.model,
weight_decay=self.args.l2_lambda,
p=2)
# @override
def save_ckpt(self):
checkpoint = {
'model_state_dict': self.model.state_dict(),
'best_dev_performance': self.best_dev_performance
}
model_name = os.path.join(self.args.checkpoint_dir,
"{}.ckpt".format(self.args.experiment_name))
torch.save(checkpoint, model_name)
print("Save model as %s" % model_name)
def train(self, start_epoch, end_epoch):
eval_every = self.args.eval_every
print("Strat Training------------------")
self.show_norm()
if self.args.norm_one:
self.model.norm_one()
elif self.args.norm_emb:
self.model.norm_emb()
if self.args.norm_user:
self.model.norm_user()
#self.evaluator_kg.evaluate(self.kg_eval_loader, "EVAL")
#self.evaluator_kg.evaluate(self.kg_test_loader, "TEST")
# if self.args.need_pretrain:
# self.pretrain(start_epoch, end_epoch)
for epoch in range(start_epoch, end_epoch + 1):
st = time.time()
train_loss, pos_loss, neg_loss = self.train_epoch_kg()
self.show_norm()
print("Epoch: {}, loss: {:.4f}, pos: {:.4f},"
" neg: {:.4f}, time: {}".format(epoch + 1, train_loss,
pos_loss, neg_loss,
time.time() - st))
if (epoch + 1) % eval_every == 0 and epoch > 0:
mr_raw, hits_10_raw, mrr_raw,\
mr_fil, hits_10_fil, mrr_fil = self.evaluator_kg.evaluate(self.kg_eval_loader, "EVAL")
#mr_raw, hits_10_raw, mrr_raw, mr_fil, hits_10_fil, mrr_fil = self.evaluate(self.kg_eval_loader, "EVAL")
if mrr_fil > self.best_dev_performance:
self.best_dev_performance = mrr_fil
self.save_ckpt()
self.reset_time = 0
# To guarantee the correctness of evaluation
else:
self.logger.info(
'No improvement after one evaluation iter.')
self.reset_time += 1
if self.reset_time >= 5:
self.logger.info(
'No improvement after 5 evaluation. Early Stopping.')
break
self.logger.info('Train Done! Evaluate on testset with saved model')
print("End Training------------------")
self.evaluate_best()
def dis_step(self, heads, rels, tails, neg_tails, sample_weight, user_rep):
self.optim.zero_grad()
query_now, norm_q = self.model.form_query(heads, rels, user_rep)
preds, norm_p = self.model.query_judge(query_now, tails)
preds_neg, norm_n = self.model.query_judge(query_now, neg_tails)
random_ratio = self.args.label_smoothing_epsilon / self.args.n_sample
answers_true = torch.ones_like(preds) * (
1.0 - self.args.label_smoothing_epsilon)
answers_false = torch.zeros_like(preds_neg) + random_ratio
loss_pos = self.loss_func_kg(preds, answers_true)
loss_pos = (loss_pos * sample_weight).sum() # / sample_weight.sum()
loss_neg = torch.sum(self.loss_func_kg(preds_neg, answers_false),
dim=1)
loss_neg = (loss_neg * sample_weight).sum() # / sample_weight.sum()
# loss_reg = self.loss_reg(self.model)
loss = loss_pos + loss_neg
# losses_reg.append(loss_reg.item())
loss.backward()
torch.nn.utils.clip_grad_norm_(
[param for name, param in self.model.named_parameters()],
self.args.clipping_max_value)
self.optim.step()
if self.args.norm_one:
self.model.norm_one()
elif self.args.norm_emb:
self.model.norm_emb()
if self.args.norm_user:
self.model.norm_user()
return loss.item(), loss_pos.item(), loss_neg.item()
def train_epoch_kg(self):
path_list = self.sample_epoch_edges()
self.model.train()
losses = []
losses_pos = []
losses_neg = []
step = 0
path_reverse = self.sample_epoch_reverse()
rgcn_kernel = self.get_rgcn_kenel(path_reserve=path_reverse,
user_set=set(range(self.user_total)))
for heads, rels, tails, neg_tails, sample_weight in self.kg_train_loader:
step += 1
self.optim.zero_grad()
max_hop, ent_hop, hop_map, unreach_ents, new_order = self.batch_reserve_ents(
heads, path_list)
rs_graph, kg_graph = make_kernel_batch(path_list, max_hop, ent_hop,
hop_map)
batch_order = torch.LongTensor(new_order).to(self.device)
user_rep = self.model.fetch_user_batch(rgcn_kernel=rgcn_kernel,
rs_graph=rs_graph,
kg_graph=kg_graph,
tp_hop=max_hop,
unreach_ids=unreach_ents,
batch_order=batch_order,
pretrain=False)
heads = torch.LongTensor(heads).to(self.device)
rels = torch.LongTensor(rels).to(self.device)
tails = torch.LongTensor(tails).to(self.device)
neg_tails = torch.LongTensor(neg_tails).to(self.device)
sample_weight = sample_weight.float().to(self.device).squeeze()
loss, loss_pos, loss_neg = self.dis_step(heads, rels, tails,
neg_tails, sample_weight,
user_rep)
losses.append(loss)
losses_pos.append(loss_pos)
losses_neg.append(loss_neg)
if self.decay_rate:
self.scheduler.step()
mean_losses = np.mean(losses)
pos_loss = np.mean(losses_pos)
neg_loss = np.mean(losses_neg)
return mean_losses, pos_loss, neg_loss
class TrackerSiamFC(Tracker):
def __init__(self, net_path=None, **kwargs):
super(TrackerSiamFC, self).__init__('SiamFC', True)
self.cfg = self.parse_args(**kwargs)
self.update_id = 2
# setup GPU device if available
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
# setup model 모델 설정하기
self.net = Net(backbone=AlexNetV1(), head=SiamFC(self.cfg.out_scale))
ops.init_weights(self.net)
# load checkpoint if provided
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
self.net = self.net.to(self.device)
# setup criterion
self.criterion = BalancedLoss()
# setup optimizer
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.cfg.initial_lr,
weight_decay=self.cfg.weight_decay,
momentum=self.cfg.momentum)
# setup lr scheduler
gamma = np.power(self.cfg.ultimate_lr / self.cfg.initial_lr,
1.0 / self.cfg.epoch_num)
self.lr_scheduler = ExponentialLR(self.optimizer, gamma)
def parse_args(self, **kwargs):
# default parameters
cfg = {
# basic parameters
'out_scale': 0.001,
'exemplar_sz': 127,
'instance_sz': 255,
'context': 0.5,
# inference parameters
'scale_num': 9,
'scale_num_res': 3,
'scale_step': 1.0375,
'scale_lr': 0.59,
'scale_penalty': 0.9745,
'window_influence': 0.176,
'response_sz': 17, #17 x 17
'response_up': 16,
'total_stride': 8,
# train parameters
'epoch_num': 50,
'batch_size': 8,
'num_workers': 32,
'initial_lr': 1e-2,
'ultimate_lr': 1e-5,
'weight_decay': 5e-4,
'momentum': 0.9,
'r_pos': 16,
'r_neg': 0
}
for key, val in kwargs.items():
if key in cfg:
cfg.update({key: val})
return namedtuple('Config', cfg.keys())(**cfg)
@torch.no_grad()
def init(self, img, box):
# set to evaluation mode
self.net.eval()
# convert box to 0-indexed and center based [y, x, h, w]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
self.center, self.target_sz = box[:2], box[2:]
# create hanning window
self.upscale_sz = self.cfg.response_up * self.cfg.response_sz
self.hann_window = np.outer(
np.hanning(
self.upscale_sz
), #np.hanning() https://numpy.org/doc/stable/reference/generated/numpy.hanning.html
np.hanning(self.upscale_sz))
self.hann_window /= self.hann_window.sum()
# search scale factors
self.scale_factors = self.cfg.scale_step**np.linspace(
-(self.cfg.scale_num // 8), self.cfg.scale_num // 8,
self.cfg.scale_num)
# search scale factors
self.scale_factors_res = self.cfg.scale_step**np.linspace(
-(self.cfg.scale_num_res // 2), self.cfg.scale_num_res // 2,
self.cfg.scale_num_res)
# exemplar and search sizes
context = self.cfg.context * np.sum(self.target_sz)
self.z_sz = np.sqrt(np.prod(self.target_sz + context))
self.x_sz = self.z_sz * self.cfg.instance_sz / self.cfg.exemplar_sz
# exemplar image
self.avg_color = np.mean(img, axis=(0, 1))
z = ops.crop_and_resize(img,
self.center,
self.z_sz,
out_size=self.cfg.exemplar_sz,
border_value=self.avg_color)
# exemplar features
z = torch.from_numpy(z).to(self.device).permute(
2, 0, 1).unsqueeze(0).float()
self.kernel = self.net.backbone(z) #AlexNetV1()
@torch.no_grad()
def update(self, count, img_files):
# set to evaluation mode
self.net.eval()
x = []
#setting scale_num and scale_factors
if count == 1:
for fr, img_file in enumerate(img_files):
img = ops.read_image(img_file)
x.append(
ops.crop_and_resize(img,
self.center,
self.x_sz * 1.0,
out_size=self.cfg.instance_sz,
border_value=self.avg_color)) #0.7
else:
for fr, img_file in enumerate(img_files):
img = ops.read_image(img_file)
y = [
ops.crop_and_resize(img,
self.center,
self.x_sz * f,
out_size=self.cfg.instance_sz,
border_value=self.avg_color)
for f in self.scale_factors_res
]
x.extend(y)
# search images
search = x
x = np.stack(x, axis=0)
x = torch.from_numpy(x).to(self.device).permute(0, 3, 1, 2).float()
# responses
x = self.net.backbone(x) #AlexNetV1()
responses = self.net.head(self.kernel, x)
responses = responses.squeeze(1).cpu().numpy()
scale_num = self.cfg.scale_num
scale_factors = self.scale_factors
panalty_score = 8
# upsample responses and penalize scale changes
responses = np.stack([
cv2.resize(u, (self.upscale_sz, self.upscale_sz),
interpolation=cv2.INTER_CUBIC) for u in responses
])
responses[:scale_num // panalty_score] *= self.cfg.scale_penalty
responses[scale_num // panalty_score + 1:] *= self.cfg.scale_penalty
# peak scale
pc = np.amax(responses, axis=(1, 2))
scale_id = np.argmax(np.amax(responses, axis=(1, 2)))
if count == 1:
file_id = scale_id
else:
if scale_id < 3:
file_id = 0
elif scale_id > 2 and scale_id < 6:
file_id = 1
else:
file_id = 2
png = img_files[file_id].split('\\')[-1]
png_id = int(png.split('.')[-2])
# peak location
response = responses[scale_id]
response -= response.min()
response /= response.sum() + 1e-16
response = (1 - self.cfg.window_influence
) * response + self.cfg.window_influence * self.hann_window
loc = np.unravel_index(response.argmax(), response.shape)
# locate target center
disp_in_response = np.array(loc) - (self.upscale_sz - 1) / 2
disp_in_instance = disp_in_response * self.cfg.total_stride / self.cfg.response_up
disp_in_image = disp_in_instance * self.x_sz * scale_factors[
scale_id] / self.cfg.instance_sz
self.center += disp_in_image
scale = (1 - self.cfg.scale_lr
) * 1.0 + self.cfg.scale_lr * scale_factors[scale_id]
self.target_sz *= scale
self.z_sz *= scale
self.x_sz *= scale
# return 1-indexed and left-top based bounding box
box = np.array([
self.center[1] + 1 - (self.target_sz[1] - 1) / 2,
self.center[0] + 1 - (self.target_sz[0] - 1) / 2,
self.target_sz[1], self.target_sz[0]
])
return box, responses, search, png_id
def track(self, first_img, focal_dirs, show_imgs, box, visualize=False):
frame_num = len(focal_dirs)
boxes = np.zeros((frame_num, 4))
boxes[0] = box
times = np.zeros(frame_num)
response = np.zeros((3, 272, 272))
search = []
picked = []
gt = open("recode/bbox.txt", 'w')
try:
for f, focal_dir in enumerate(focal_dirs):
focal_plane = glob.glob(focal_dir + "/*")
begin = time.time()
if f == 0:
self.init(first_img, box)
elif f == 1:
sharp = sharpfiles(
focal_dir) #focal_plane[27] #처음은 027.png
boxes[f, :], response, search, scale_id = self.update(
f, sharp) #sharp
else:
focal_planes = focal_plane[scale_id - 1:scale_id + 1 + 1]
boxes[f, :], response, search, scale_id = self.update(
f, focal_planes) #focal 이미지 리스트
times[f] = time.time() - begin
if response[0][0][0] == 0.0:
visualize = False
else:
visualize = True
if visualize == True:
dimg = ops.read_image(show_imgs[f])
save_img, centerloc = ops.show_image(dimg,
boxes[f, :],
None,
num_name=f)
save_img = cv2.putText(save_img, str(scale_id), (100, 30),
cv2.FONT_HERSHEY_SIMPLEX, 0.5,
(0, 255, 0), 2)
savepath = "./image/{0:0=3d}.png".format(f)
cv2.imwrite(savepath, save_img)
#groundtruth
center_coo = "%d,%d,%d,%d\n" % (int(
centerloc[0]), int(centerloc[1]), int(
centerloc[2]), int(centerloc[3]))
gt.write(center_coo)
finally:
gt.close()
return boxes, times, response
def train_step(self, batch, backward=True):
# set network mode
self.net.train(backward)
# parse batch data
z = batch[0].to(self.device, non_blocking=self.cuda)
x = batch[1].to(self.device, non_blocking=self.cuda)
with torch.set_grad_enabled(backward):
# inference
responses = self.net(z, x)
# calculate loss
labels = self._create_labels(responses.size())
loss = self.criterion(responses, labels)
if backward:
# back propagation
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
@torch.enable_grad()
def train_over(self, seqs, val_seqs=None, save_dir='pretrained'):
# set to train mode
self.net.train()
# create save_dir folder
if not os.path.exists(save_dir):
os.makedirs(save_dir)
# setup dataset
transforms = SiamFCTransforms(exemplar_sz=self.cfg.exemplar_sz,
instance_sz=self.cfg.instance_sz,
context=self.cfg.context)
dataset = Pair(seqs=seqs, transforms=transforms)
# setup dataloader
dataloader = DataLoader(dataset,
batch_size=self.cfg.batch_size,
shuffle=True,
num_workers=self.cfg.num_workers,
pin_memory=self.cuda,
drop_last=True)
# loop over epochs
for epoch in range(self.cfg.epoch_num):
# update lr at each epoch
self.lr_scheduler.step(epoch=epoch)
# loop over dataloader
for it, batch in enumerate(dataloader):
loss = self.train_step(batch, backward=True)
print('Epoch: {} [{}/{}] Loss: {:.5f}'.format(
epoch + 1, it + 1, len(dataloader), loss))
sys.stdout.flush()
# save checkpoint
if not os.path.exists(save_dir):
os.makedirs(save_dir)
net_path = os.path.join(save_dir,
'siamfc_alexnet_e%d.pth' % (epoch + 1))
torch.save(self.net.state_dict(), net_path)
def _create_labels(self, size):
# skip if same sized labels already created
if hasattr(self, 'labels') and self.labels.size() == size:
return self.labels
def logistic_labels(x, y, r_pos, r_neg):
dist = np.abs(x) + np.abs(y) # block distance
labels = np.where(
dist <= r_pos, np.ones_like(x),
np.where(dist < r_neg,
np.ones_like(x) * 0.5, np.zeros_like(x)))
return labels
# distances along x- and y-axis
n, c, h, w = size
x = np.arange(w) - (w - 1) / 2
y = np.arange(h) - (h - 1) / 2
x, y = np.meshgrid(x, y)
# create logistic labels
r_pos = self.cfg.r_pos / self.cfg.total_stride
r_neg = self.cfg.r_neg / self.cfg.total_stride
labels = logistic_labels(x, y, r_pos, r_neg)
# repeat to size
labels = labels.reshape((1, 1, h, w))
labels = np.tile(labels, (n, c, 1, 1))
# convert to tensors
self.labels = torch.from_numpy(labels).to(self.device).float()
return self.labels
class PPO:
def __init__(
self,
seed: int,
gamma: float,
tau: float,
clip_param: float,
vf_c: float,
ent_c: float,
input_shape: int,
hidden_units_value: list,
hidden_units_actor: list,
batch_size: int,
lr: float,
activation: str,
optimizer_name: str,
batch_norm_input: bool,
batch_norm_value_out: bool,
action_space,
policy_type: str,
init_pol_std: float,
min_pol_std: float,
beta_1: float = 0.9,
beta_2: float = 0.999,
eps_opt: float = 1e-07,
lr_schedule: Optional[str] = None,
exp_decay_rate: Optional[float] = None,
step_size: Optional[int] = None,
std_transform: str = "softplus",
init_last_layers: str = "rescaled",
rng=None,
modelname: str = "PPO act_crt",
):
if rng is not None:
self.rng = rng
else:
self.rng = np.random.RandomState(seed)
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# self.device = torch.device('cpu')
self.gamma = gamma
self.tau = tau
self.clip_param = clip_param
self.vf_c = vf_c
self.ent_c = ent_c
self.batch_size = batch_size
self.beta_1 = beta_1
self.eps_opt = eps_opt
self.action_space = action_space
self.policy_type = policy_type
self.num_actions = self.action_space.get_n_actions(policy_type=self.policy_type)
self.batch_norm_input = batch_norm_input
self.experience = {
"state": [],
"action": [],
"reward": [],
"log_prob": [],
"value": [],
"returns": [],
"advantage": [],
}
self.model = PPOActorCritic(
seed,
input_shape,
activation,
hidden_units_value,
hidden_units_actor,
self.num_actions,
batch_norm_input,
batch_norm_value_out,
self.policy_type,
init_pol_std,
min_pol_std,
std_transform,
init_last_layers,
modelname,
)
self.model.to(self.device)
self.optimizer_name = optimizer_name
if optimizer_name == "adam":
self.optimizer = optim.Adam(
self.model.parameters(),
lr=lr,
betas=(beta_1, beta_2),
eps=eps_opt,
weight_decay=0,
amsgrad=False,
)
elif optimizer_name == "rmsprop":
self.optimizer = optim.RMSprop(
self.model.parameters(),
lr=lr,
alpha=beta_1,
eps=eps_opt,
weight_decay=0,
momentum=0,
centered=False,
)
if lr_schedule == "step":
self.scheduler = StepLR(
optimizer=self.optimizer, step_size=step_size, gamma=exp_decay_rate
)
elif lr_schedule == "exponential":
self.scheduler = ExponentialLR(
optimizer=self.optimizer, gamma=exp_decay_rate
)
else:
self.scheduler = None
if self.policy_type == 'continuous':
self.std_hist = []
self.entropy_hist = []
elif self.policy_type == 'discrete':
self.logits_hist = []
self.entropy_hist = []
def train(self, state, action, old_log_probs, return_, advantage):
advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-5)
self.model.train()
dist, value = self.model(state)
entropy = dist.entropy().mean()
if self.policy_type == 'continuous':
new_log_probs = dist.log_prob(action)
# self.std_hist.append(self.model.log_std.exp().detach().cpu().numpy().ravel())
# self.entropy_hist.append(entropy.detach().cpu().numpy().ravel())
elif self.policy_type == 'discrete':
new_log_probs = dist.log_prob(action.reshape(-1)).reshape(-1,1)
# self.logits_hist.append(dist.logits.detach().cpu().numpy())
# self.entropy_hist.append(entropy.detach().cpu().numpy().ravel())
ratio = (new_log_probs - old_log_probs).exp() # log properties
surr1 = ratio * advantage
surr2 = (
torch.clamp(ratio, 1.0 / (1 + self.clip_param), 1.0 + self.clip_param)
* advantage
)
actor_loss = -torch.min(surr1, surr2).mean()
critic_loss = (return_ - value).pow(2).mean()
# the loss is negated in order to be maximized
self.loss = self.vf_c * critic_loss + actor_loss - self.ent_c * entropy
self.optimizer.zero_grad()
self.loss.backward()
self.optimizer.step()
if self.scheduler:
self.scheduler.step()
def act(self, states):
# useful when the states are single dimensional
self.model.eval()
# make 1D tensor to 2D
with torch.no_grad():
states = torch.from_numpy(states).float().unsqueeze(0)
states = states.to(self.device)
return self.model(states)
def compute_gae(self, next_value, recompute_value=False):
if recompute_value:
self.model.eval()
with torch.no_grad():
_, values = self.model(
torch.Tensor(self.experience["state"]).to(self.device)
)
self.experience["value"] = [
np.array(v, dtype=float) for v in values.detach().cpu().tolist()
]
# for i in range(len(self.experience["value"])):
# _, value = self.act(self.experience["state"][i])
# self.experience["value"][i] = value.detach().cpu().numpy().ravel()
rewards = self.experience["reward"]
values = self.experience["value"]
values = values + [next_value]
gae = 0
returns = []
for step in reversed(range(len(rewards))):
delta = rewards[step] + self.gamma * values[step + 1] - values[step]
gae = delta + self.gamma * self.tau * gae
returns.insert(0, gae + values[step])
# add estimated returns and advantages to the experience
self.experience["returns"] = returns
advantage = [returns[i] - values[i] for i in range(len(returns))]
self.experience["advantage"] = advantage
# add way to reset experience after one rollout
def add_experience(self, exp):
for key, value in exp.items():
self.experience[key].append(value)
def reset_experience(self):
self.experience = {
"state": [],
"action": [],
"reward": [],
"log_prob": [],
"value": [],
"returns": [],
"advantage": [],
}
def ppo_iter(self):
# pick a batch from the rollout
states = np.asarray(self.experience["state"])
actions = np.asarray(self.experience["action"])
log_probs = np.asarray(self.experience["log_prob"])
returns = np.asarray(self.experience["returns"])
advantage = np.asarray(self.experience["advantage"])
len_rollout = states.shape[0]
ids = self.rng.permutation(len_rollout)
ids = np.array_split(ids, len_rollout // self.batch_size)
for i in range(len(ids)):
yield (
torch.from_numpy(states[ids[i], :]).float().to(self.device),
torch.from_numpy(actions[ids[i], :]).float().to(self.device),
torch.from_numpy(log_probs[ids[i], :]).float().to(self.device),
torch.from_numpy(returns[ids[i], :]).float().to(self.device),
torch.from_numpy(advantage[ids[i], :]).float().to(self.device),
)
def getBack(self, var_grad_fn):
print(var_grad_fn)
for n in var_grad_fn.next_functions:
if n[0]:
try:
tensor = getattr(n[0], "variable")
print(n[0])
print("Tensor with grad found:", tensor)
print(" - gradient:", tensor.grad)
print()
except AttributeError as e:
self.getBack(n[0])
def save_diagnostics(self,path):
if self.policy_type == 'continuous':
np.save(os.path.join(path, "std_hist"), np.array(self.std_hist))
np.save(os.path.join(path, "entropy_hist"), np.array(self.entropy_hist))
elif self.policy_type == 'discrete':
np.save(os.path.join(path, "logits_hist"), np.array(self.logits_hist, dtype=object))
np.save(os.path.join(path, "entropy_hist"), np.array(self.entropy_hist))
def main() -> None:
"""Entrypoint.
"""
config: Any = importlib.import_module(args.config)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_data = tx.data.MonoTextData(config.train_data_hparams, device=device)
val_data = tx.data.MonoTextData(config.val_data_hparams, device=device)
test_data = tx.data.MonoTextData(config.test_data_hparams, device=device)
iterator = tx.data.DataIterator({
"train": train_data,
"valid": val_data,
"test": test_data
})
opt_vars = {
'learning_rate': config.lr_decay_hparams["init_lr"],
'best_valid_nll': 1e100,
'steps_not_improved': 0,
'kl_weight': config.kl_anneal_hparams["start"]
}
decay_cnt = 0
max_decay = config.lr_decay_hparams["max_decay"]
decay_factor = config.lr_decay_hparams["decay_factor"]
decay_ts = config.lr_decay_hparams["threshold"]
save_dir = f"./models/{config.dataset}"
if not os.path.exists(save_dir):
os.makedirs(save_dir)
suffix = f"{config.dataset}_{config.decoder_type}Decoder.ckpt"
save_path = os.path.join(save_dir, suffix)
# KL term annealing rate
anneal_r = 1.0 / (config.kl_anneal_hparams["warm_up"] *
(len(train_data) / config.batch_size))
vocab = train_data.vocab
model = VAE(train_data.vocab.size, config)
model.to(device)
start_tokens = torch.full((config.batch_size, ),
vocab.bos_token_id,
dtype=torch.long).to(device)
end_token = vocab.eos_token_id
optimizer = tx.core.get_optimizer(params=model.parameters(),
hparams=config.opt_hparams)
scheduler = ExponentialLR(optimizer, decay_factor)
def _run_epoch(epoch: int, mode: str, display: int = 10) \
-> Tuple[Tensor, float]:
iterator.switch_to_dataset(mode)
if mode == 'train':
model.train()
opt_vars["kl_weight"] = min(1.0, opt_vars["kl_weight"] + anneal_r)
kl_weight = opt_vars["kl_weight"]
else:
model.eval()
kl_weight = 1.0
step = 0
start_time = time.time()
num_words = 0
nll_total = 0.
avg_rec = tx.utils.AverageRecorder()
for batch in iterator:
ret = model(batch, kl_weight, start_tokens, end_token)
if mode == "train":
opt_vars["kl_weight"] = min(1.0,
opt_vars["kl_weight"] + anneal_r)
kl_weight = opt_vars["kl_weight"]
ret["nll"].backward()
optimizer.step()
optimizer.zero_grad()
batch_size = len(ret["lengths"])
num_words += torch.sum(ret["lengths"]).item()
nll_total += ret["nll"].item() * batch_size
avg_rec.add([
ret["nll"].item(), ret["kl_loss"].item(),
ret["rc_loss"].item()
], batch_size)
if step % display == 0 and mode == 'train':
nll = avg_rec.avg(0)
klw = opt_vars["kl_weight"]
KL = avg_rec.avg(1)
rc = avg_rec.avg(2)
log_ppl = nll_total / num_words
ppl = math.exp(log_ppl)
time_cost = time.time() - start_time
print(
f"{mode}: epoch {epoch}, step {step}, nll {nll:.4f}, "
f"klw {klw:.4f}, KL {KL:.4f}, rc {rc:.4f}, "
f"log_ppl {log_ppl:.4f}, ppl {ppl:.4f}, "
f"time_cost {time_cost:.1f}",
flush=True)
step += 1
nll = avg_rec.avg(0)
KL = avg_rec.avg(1)
rc = avg_rec.avg(2)
log_ppl = nll_total / num_words
ppl = math.exp(log_ppl)
print(f"\n{mode}: epoch {epoch}, nll {nll:.4f}, KL {KL:.4f}, "
f"rc {rc:.4f}, log_ppl {log_ppl:.4f}, ppl {ppl:.4f}")
return nll, ppl # type: ignore
@torch.no_grad()
def _generate(start_tokens: torch.LongTensor,
end_token: int,
filename: Optional[str] = None):
ckpt = torch.load(args.model)
model.load_state_dict(ckpt['model'])
model.eval()
batch_size = train_data.batch_size
dst = MultivariateNormalDiag(loc=torch.zeros(batch_size,
config.latent_dims),
scale_diag=torch.ones(
batch_size, config.latent_dims))
latent_z = dst.rsample().to(device)
helper = model.decoder.create_helper(decoding_strategy='infer_sample',
start_tokens=start_tokens,
end_token=end_token)
outputs = model.decode(helper=helper,
latent_z=latent_z,
max_decoding_length=100)
sample_tokens = vocab.map_ids_to_tokens_py(outputs.sample_id.cpu())
if filename is None:
fh = sys.stdout
else:
fh = open(filename, 'w', encoding='utf-8')
for sent in sample_tokens:
sent = tx.utils.compat_as_text(list(sent))
end_id = len(sent)
if vocab.eos_token in sent:
end_id = sent.index(vocab.eos_token)
fh.write(' '.join(sent[:end_id + 1]) + '\n')
print('Output done')
fh.close()
if args.mode == "predict":
_generate(start_tokens, end_token, args.out)
return
# Counts trainable parameters
total_parameters = sum(param.numel() for param in model.parameters())
print(f"{total_parameters} total parameters")
best_nll = best_ppl = 0.
for epoch in range(config.num_epochs):
_, _ = _run_epoch(epoch, 'train', display=200)
val_nll, _ = _run_epoch(epoch, 'valid')
test_nll, test_ppl = _run_epoch(epoch, 'test')
if val_nll < opt_vars['best_valid_nll']:
opt_vars['best_valid_nll'] = val_nll
opt_vars['steps_not_improved'] = 0
best_nll = test_nll
best_ppl = test_ppl
states = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict(),
"scheduler": scheduler.state_dict()
}
torch.save(states, save_path)
else:
opt_vars['steps_not_improved'] += 1
if opt_vars['steps_not_improved'] == decay_ts:
old_lr = opt_vars['learning_rate']
opt_vars['learning_rate'] *= decay_factor
opt_vars['steps_not_improved'] = 0
new_lr = opt_vars['learning_rate']
ckpt = torch.load(save_path)
model.load_state_dict(ckpt['model'])
optimizer.load_state_dict(ckpt['optimizer'])
scheduler.load_state_dict(ckpt['scheduler'])
scheduler.step()
print(f"-----\nchange lr, old lr: {old_lr}, "
f"new lr: {new_lr}\n-----")
decay_cnt += 1
if decay_cnt == max_decay:
break
print(f"\nbest testing nll: {best_nll:.4f},"
f"best testing ppl {best_ppl:.4f}\n")
loss_kld=loss_kld_mean.measure,
epoch=i + 1)
if slurm:
progress.update(
progress.value,
loss_nle=loss_nle_mean.measure,
loss_encoder=loss_encoder_mean.measure,
loss_decoder=loss_decoder_mean.measure,
loss_discriminator=loss_discriminator_mean.measure,
loss_mse_layer=loss_reconstruction_layer_mean.measure,
loss_kld=loss_kld_mean.measure,
epoch=i + 1)
# EPOCH END
lr_encoder.step()
lr_decoder.step()
lr_discriminator.step()
margin *= decay_margin
equilibrium *= decay_equilibrium
torch.save({
'epoch': step_index,
"net": net.module.state_dict()
}, save_path + f'model_epoch_{step_index}.tar')
#margin non puo essere piu alto di equilibrium
if margin > equilibrium:
equilibrium = margin
lambda_mse *= decay_mse
if lambda_mse > 1:
lambda_mse = 1
progress.finish()
class SamplingMultitaskTrainer:
def __init__(self,
dataset=None,
model_name=None,
model_params=None,
trainer_params=None,
restore=None,
device=None,
pretrained_embeddings_path=None,
tokenizer_path=None):
self.graph_model = model_name(dataset.g, **model_params).to(device)
self.model_params = model_params
self.trainer_params = trainer_params
self.device = device
self.epoch = 0
self.batch = 0
self.dtype = torch.float32
self.create_node_embedder(
dataset,
tokenizer_path,
n_dims=model_params["h_dim"],
pretrained_path=pretrained_embeddings_path,
n_buckets=trainer_params["embedding_table_size"])
self.summary_writer = SummaryWriter(self.model_base_path)
self.ee_node_name = ElementEmbedderWithBpeSubwords(
elements=dataset.load_node_names(),
nodes=dataset.nodes,
emb_size=self.elem_emb_size,
tokenizer_path=tokenizer_path).to(self.device)
self.ee_var_use = ElementEmbedderWithBpeSubwords(
elements=dataset.load_var_use(),
nodes=dataset.nodes,
emb_size=self.elem_emb_size,
tokenizer_path=tokenizer_path).to(self.device)
self.ee_api_call = ElementEmbedderBase(
elements=dataset.load_api_call(),
nodes=dataset.nodes,
compact_dst=False,
dst_to_global=True)
self.lp_node_name = LinkPredictor(self.ee_node_name.emb_size +
self.graph_model.emb_size).to(
self.device)
self.lp_var_use = LinkPredictor(self.ee_var_use.emb_size +
self.graph_model.emb_size).to(
self.device)
self.lp_api_call = LinkPredictor(self.graph_model.emb_size +
self.graph_model.emb_size).to(
self.device)
if restore:
self.restore_from_checkpoint(self.model_base_path)
self.optimizer = self._create_optimizer()
self.lr_scheduler = ExponentialLR(self.optimizer, gamma=1.0)
self.best_score = BestScoreTracker()
self._create_loaders(*self._get_training_targets())
def create_node_embedder(self,
dataset,
tokenizer_path,
n_dims=None,
pretrained_path=None,
n_buckets=500000):
from SourceCodeTools.nlp.embed.fasttext import load_w2v_map
if pretrained_path is not None:
pretrained = load_w2v_map(pretrained_path)
else:
pretrained = None
if pretrained_path is None and n_dims is None:
raise ValueError(
f"Specify embedding dimensionality or provide pretrained embeddings"
)
elif pretrained_path is not None and n_dims is not None:
assert n_dims == pretrained.n_dims, f"Requested embedding size and pretrained embedding " \
f"size should match: {n_dims} != {pretrained.n_dims}"
elif pretrained_path is not None and n_dims is None:
n_dims = pretrained.n_dims
if pretrained is not None:
logging.info(f"Loading pretrained embeddings...")
logging.info(f"Input embedding size is {n_dims}")
self.node_embedder = NodeEmbedder(
nodes=dataset.nodes,
emb_size=n_dims,
# tokenizer_path=tokenizer_path,
dtype=self.dtype,
pretrained=dataset.buckets_from_pretrained_embeddings(
pretrained_path, n_buckets)
if pretrained_path is not None else None,
n_buckets=n_buckets)
# self.node_embedder(node_type="node_", node_ids=torch.LongTensor([0]))
# self.node_embedder(node_type="node_", node_ids=torch.LongTensor([13749]))
# self.node_embedder(node_type="node_", node_ids=torch.LongTensor([13754]))
# node_, 0 matplotlib
# node_ 13749 Renderer
# node_ 13754 ▁renderer
# print()
@property
def lr(self):
return self.trainer_params['lr']
@property
def batch_size(self):
return self.trainer_params['batch_size']
@property
def sampling_neighbourhood_size(self):
return self.trainer_params['sampling_neighbourhood_size']
@property
def neg_sampling_factor(self):
return self.trainer_params['neg_sampling_factor']
@property
def epochs(self):
return self.trainer_params['epochs']
@property
def elem_emb_size(self):
return self.trainer_params['elem_emb_size']
@property
def node_name_file(self):
return self.trainer_params['node_name_file']
@property
def var_use_file(self):
return self.trainer_params['var_use_file']
@property
def call_seq_file(self):
return self.trainer_params['call_seq_file']
@property
def model_base_path(self):
return self.trainer_params['model_base_path']
@property
def pretraining(self):
return self.epoch >= self.trainer_params['pretraining_phase']
@property
def do_save(self):
return self.trainer_params['save_checkpoints']
# def _extract_embed(self, node_embed, input_nodes):
# emb = {}
# for node_type, nid in input_nodes.items():
# emb[node_type] = node_embed[node_type][nid]
# return emb
def write_summary(self, scores, batch_step):
# main_name = os.path.basename(self.model_base_path)
for var, val in scores.items():
# self.summary_writer.add_scalar(f"{main_name}/{var}", val, batch_step)
self.summary_writer.add_scalar(var, val, batch_step)
# self.summary_writer.add_scalars(main_name, scores, batch_step)
def write_hyperparams(self, scores, epoch):
params = copy(self.model_params)
params["epoch"] = epoch
main_name = os.path.basename(self.model_base_path)
params = {
k: v
for k, v in params.items()
if type(v) in {int, float, str, bool, torch.Tensor}
}
main_name = os.path.basename(self.model_base_path)
scores = {f"h_metric/{k}": v for k, v in scores.items()}
self.summary_writer.add_hparams(params,
scores,
run_name=f"h_metric/{epoch}")
def _extract_embed(self, input_nodes):
emb = {}
for node_type, nid in input_nodes.items():
emb[node_type] = self.node_embedder(
node_type=node_type,
node_ids=nid,
train_embeddings=self.pretraining).to(self.device)
return emb
def _logits_batch(self, input_nodes, blocks):
cumm_logits = []
if self.use_types:
# emb = self._extract_embed(self.graph_model.node_embed(), input_nodes)
emb = self._extract_embed(input_nodes)
else:
if self.ntypes is not None:
# single node type
key = next(iter(self.ntypes))
input_nodes = {key: input_nodes}
# emb = self._extract_embed(self.graph_model.node_embed(), input_nodes)
emb = self._extract_embed(input_nodes)
else:
emb = self.node_embedder(node_ids=input_nodes,
train_embeddings=self.pretraining)
# emb = self.graph_model.node_embed()[input_nodes]
logits = self.graph_model(emb, blocks)
if self.use_types:
for ntype in self.graph_model.g.ntypes:
logits_ = logits.get(ntype, None)
if logits_ is None:
continue
cumm_logits.append(logits_)
else:
if self.ntypes is not None:
# single node type
key = next(iter(self.ntypes))
logits_ = logits[key]
else:
logits_ = logits
cumm_logits.append(logits_)
return torch.cat(cumm_logits)
def seeds_to_global(self, seeds):
if type(seeds) is dict:
indices = [
self.graph_model.g.nodes[ntype].data["global_graph_id"][
seeds[ntype]] for ntype in seeds
]
return torch.cat(indices, dim=0)
else:
return seeds
def _logits_embedder(self,
node_embeddings,
elem_embedder,
link_predictor,
seeds,
negative_factor=1):
k = negative_factor
indices = self.seeds_to_global(seeds)
batch_size = len(indices)
node_embeddings_batch = node_embeddings
element_embeddings = elem_embedder(elem_embedder[indices.tolist()].to(
self.device))
positive_batch = torch.cat([node_embeddings_batch, element_embeddings],
1)
labels_pos = torch.ones(batch_size, dtype=torch.long)
node_embeddings_neg_batch = node_embeddings_batch.repeat(k, 1)
negative_random = elem_embedder(
elem_embedder.sample_negative(batch_size * k).to(self.device))
negative_batch = torch.cat(
[node_embeddings_neg_batch, negative_random], 1)
labels_neg = torch.zeros(batch_size * k, dtype=torch.long)
batch = torch.cat([positive_batch, negative_batch], 0)
labels = torch.cat([labels_pos, labels_neg], 0).to(self.device)
logits = link_predictor(batch)
return logits, labels
def _handle_non_unique(self, non_unique_ids):
id_list = non_unique_ids.tolist()
unique_ids = list(set(id_list))
new_position = dict(zip(unique_ids, range(len(unique_ids))))
slice_map = torch.tensor(list(map(lambda x: new_position[x], id_list)),
dtype=torch.long)
return torch.tensor(unique_ids, dtype=torch.long), slice_map
def _logits_nodes(self,
node_embeddings,
elem_embedder,
link_predictor,
create_dataloader,
src_seeds,
negative_factor=1):
k = negative_factor
indices = self.seeds_to_global(src_seeds)
batch_size = len(indices)
node_embeddings_batch = node_embeddings
next_call_indices = elem_embedder[
indices.tolist()] # this assumes indices is torch tensor
# dst targets are not unique
unique_dst, slice_map = self._handle_non_unique(next_call_indices)
assert unique_dst[slice_map].tolist() == next_call_indices.tolist()
dataloader = create_dataloader(unique_dst)
input_nodes, dst_seeds, blocks = next(iter(dataloader))
blocks = [blk.to(self.device) for blk in blocks]
assert dst_seeds.shape == unique_dst.shape
assert dst_seeds.tolist() == unique_dst.tolist()
unique_dst_embeddings = self._logits_batch(
input_nodes, blocks) # use_types, ntypes)
next_call_embeddings = unique_dst_embeddings[slice_map.to(self.device)]
positive_batch = torch.cat(
[node_embeddings_batch, next_call_embeddings], 1)
labels_pos = torch.ones(batch_size, dtype=torch.long)
node_embeddings_neg_batch = node_embeddings_batch.repeat(k, 1)
negative_indices = torch.tensor(
elem_embedder.sample_negative(batch_size * k), dtype=torch.long
) # embeddings are sampled from 3/4 unigram distribution
unique_negative, slice_map = self._handle_non_unique(negative_indices)
assert unique_negative[slice_map].tolist() == negative_indices.tolist()
dataloader = create_dataloader(unique_negative)
input_nodes, dst_seeds, blocks = next(iter(dataloader))
blocks = [blk.to(self.device) for blk in blocks]
assert dst_seeds.shape == unique_negative.shape
assert dst_seeds.tolist() == unique_negative.tolist()
unique_negative_random = self._logits_batch(
input_nodes, blocks) # use_types, ntypes)
negative_random = unique_negative_random[slice_map.to(self.device)]
negative_batch = torch.cat(
[node_embeddings_neg_batch, negative_random], 1)
labels_neg = torch.zeros(batch_size * k, dtype=torch.long)
batch = torch.cat([positive_batch, negative_batch], 0)
labels = torch.cat([labels_pos, labels_neg], 0).to(self.device)
logits = link_predictor(batch)
return logits, labels
def _logits_node_name(self, input_nodes, seeds, blocks):
src_embs = self._logits_batch(input_nodes, blocks)
logits, labels = self._logits_embedder(
src_embs,
self.ee_node_name,
self.lp_node_name,
seeds,
negative_factor=self.neg_sampling_factor)
return logits, labels
def _logits_var_use(self, input_nodes, seeds, blocks):
src_embs = self._logits_batch(input_nodes, blocks)
logits, labels = self._logits_embedder(
src_embs,
self.ee_var_use,
self.lp_var_use,
seeds,
negative_factor=self.neg_sampling_factor)
return logits, labels
def _logits_api_call(self, input_nodes, seeds, blocks):
src_embs = self._logits_batch(input_nodes, blocks)
logits, labels = self._logits_nodes(
src_embs,
self.ee_api_call,
self.lp_api_call,
self._create_api_call_loader,
seeds,
negative_factor=self.neg_sampling_factor)
return logits, labels
def _get_training_targets(self):
if hasattr(self.graph_model.g, 'ntypes'):
self.ntypes = self.graph_model.g.ntypes
# labels = {ntype: self.graph_model.g.nodes[ntype].data['labels'] for ntype in self.ntypes}
self.use_types = True
if len(self.graph_model.g.ntypes) == 1:
# key = next(iter(labels.keys()))
# labels = labels[key]
self.use_types = False
train_idx = {
ntype: torch.nonzero(
self.graph_model.g.nodes[ntype].data['train_mask'],
as_tuple=False).squeeze()
for ntype in self.ntypes
}
val_idx = {
ntype:
torch.nonzero(self.graph_model.g.nodes[ntype].data['val_mask'],
as_tuple=False).squeeze()
for ntype in self.ntypes
}
test_idx = {
ntype: torch.nonzero(
self.graph_model.g.nodes[ntype].data['test_mask'],
as_tuple=False).squeeze()
for ntype in self.ntypes
}
else:
self.ntypes = None
# labels = g.ndata['labels']
train_idx = self.graph_model.g.ndata['train_mask']
val_idx = self.graph_model.g.ndata['val_mask']
test_idx = self.graph_model.g.ndata['test_mask']
self.use_types = False
return train_idx, val_idx, test_idx
def _evaluate_embedder(self, ee, lp, loader, neg_sampling_factor=1):
total_loss = 0
total_acc = 0
count = 0
for input_nodes, seeds, blocks in loader:
blocks = [blk.to(self.device) for blk in blocks]
src_embs = self._logits_batch(input_nodes, blocks)
logits, labels = self._logits_embedder(src_embs, ee, lp, seeds,
neg_sampling_factor)
logp = nn.functional.log_softmax(logits, 1)
loss = nn.functional.cross_entropy(logp, labels)
acc = compute_accuracy(logp.argmax(dim=1), labels)
total_loss += loss.item()
total_acc += acc
count += 1
return total_loss / count, total_acc / count
def _evaluate_nodes(self,
ee,
lp,
create_api_call_loader,
loader,
neg_sampling_factor=1):
total_loss = 0
total_acc = 0
count = 0
for input_nodes, seeds, blocks in loader:
blocks = [blk.to(self.device) for blk in blocks]
src_embs = self._logits_batch(input_nodes, blocks)
logits, labels = self._logits_nodes(src_embs, ee, lp,
create_api_call_loader, seeds,
neg_sampling_factor)
logp = nn.functional.log_softmax(logits, 1)
loss = nn.functional.cross_entropy(logp, labels)
acc = compute_accuracy(logp.argmax(dim=1), labels)
total_loss += loss.item()
total_acc += acc
count += 1
return total_loss / count, total_acc / count
def _evaluate_objectives(self, loader_node_name, loader_var_use,
loader_api_call, neg_sampling_factor):
node_name_loss, node_name_acc = self._evaluate_embedder(
self.ee_node_name,
self.lp_node_name,
loader_node_name,
neg_sampling_factor=neg_sampling_factor)
var_use_loss, var_use_acc = self._evaluate_embedder(
self.ee_var_use,
self.lp_var_use,
loader_var_use,
neg_sampling_factor=neg_sampling_factor)
api_call_loss, api_call_acc = self._evaluate_nodes(
self.ee_api_call,
self.lp_api_call,
self._create_api_call_loader,
loader_api_call,
neg_sampling_factor=neg_sampling_factor)
loss = node_name_loss + var_use_loss + api_call_loss
return loss, node_name_acc, var_use_acc, api_call_acc
def _idx_len(self, idx):
if isinstance(idx, dict):
length = 0
for key in idx:
length += len(idx[key])
else:
length = len(idx)
return length
def _get_loaders(self, train_idx, val_idx, test_idx, batch_size):
layers = self.graph_model.num_layers
# train sampler
sampler = dgl.dataloading.MultiLayerNeighborSampler(
[self.sampling_neighbourhood_size] * layers)
loader = dgl.dataloading.NodeDataLoader(self.graph_model.g,
train_idx,
sampler,
batch_size=batch_size,
shuffle=False,
num_workers=0)
# validation sampler
# we do not use full neighbor to save computation resources
val_sampler = dgl.dataloading.MultiLayerNeighborSampler(
[self.sampling_neighbourhood_size] * layers)
val_loader = dgl.dataloading.NodeDataLoader(self.graph_model.g,
val_idx,
val_sampler,
batch_size=batch_size,
shuffle=False,
num_workers=0)
# we do not use full neighbor to save computation resources
test_sampler = dgl.dataloading.MultiLayerNeighborSampler(
[self.sampling_neighbourhood_size] * layers)
test_loader = dgl.dataloading.NodeDataLoader(self.graph_model.g,
test_idx,
test_sampler,
batch_size=batch_size,
shuffle=False,
num_workers=0)
return loader, val_loader, test_loader
def _create_loaders(self, train_idx, val_idx, test_idx):
train_idx_node_name, val_idx_node_name, test_idx_node_name = self.ee_node_name.create_idx_pools(
train_idx=train_idx, val_idx=val_idx, test_idx=test_idx)
train_idx_var_use, val_idx_var_use, test_idx_var_use = self.ee_var_use.create_idx_pools(
train_idx=train_idx, val_idx=val_idx, test_idx=test_idx)
train_idx_api_call, val_idx_api_call, test_idx_api_call = self.ee_api_call.create_idx_pools(
train_idx=train_idx, val_idx=val_idx, test_idx=test_idx)
logging.info(
f"Pool sizes : train {self._idx_len(train_idx_node_name)}, "
f"val {self._idx_len(val_idx_node_name)}, "
f"test {self._idx_len(test_idx_node_name)}.")
logging.info(f"Pool sizes : train {self._idx_len(train_idx_var_use)}, "
f"val {self._idx_len(val_idx_var_use)}, "
f"test {self._idx_len(test_idx_var_use)}.")
logging.info(
f"Pool sizes : train {self._idx_len(train_idx_api_call)}, "
f"val {self._idx_len(val_idx_api_call)}, "
f"test {self._idx_len(test_idx_api_call)}.")
self.loader_node_name, self.val_loader_node_name, self.test_loader_node_name = self._get_loaders(
train_idx=train_idx_node_name,
val_idx=val_idx_node_name,
test_idx=test_idx_node_name,
batch_size=self.batch_size # batch_size_node_name
)
self.loader_var_use, self.val_loader_var_use, self.test_loader_var_use = self._get_loaders(
train_idx=train_idx_var_use,
val_idx=val_idx_var_use,
test_idx=test_idx_var_use,
batch_size=self.batch_size # batch_size_var_use
)
self.loader_api_call, self.val_loader_api_call, self.test_loader_api_call = self._get_loaders(
train_idx=train_idx_api_call,
val_idx=val_idx_api_call,
test_idx=test_idx_api_call,
batch_size=self.batch_size # batch_size_api_call
)
def _create_api_call_loader(self, indices):
sampler = dgl.dataloading.MultiLayerNeighborSampler(
[self.sampling_neighbourhood_size] * self.graph_model.num_layers)
return dgl.dataloading.NodeDataLoader(self.graph_model.g,
indices,
sampler,
batch_size=len(indices),
num_workers=0)
def _create_optimizer(self):
# optimizer = torch.optim.Adam(model.parameters(), lr=lr)
optimizer = torch.optim.Adam(
[
{
'params': self.graph_model.parameters()
},
{
'params': self.node_embedder.parameters()
},
{
'params': self.ee_node_name.parameters()
},
{
'params': self.ee_var_use.parameters()
},
# {'params': self.ee_api_call.parameters()},
{
'params': self.lp_node_name.parameters()
},
{
'params': self.lp_var_use.parameters()
},
{
'params': self.lp_api_call.parameters()
},
],
lr=self.lr)
return optimizer
def train_all(self):
"""
Training procedure for the model with node classifier
:return:
"""
for epoch in range(self.epoch, self.epochs):
self.epoch = epoch
start = time()
for i, ((input_nodes_node_name, seeds_node_name, blocks_node_name),
(input_nodes_var_use, seeds_var_use, blocks_var_use),
(input_nodes_api_call, seeds_api_call, blocks_api_call)) in \
enumerate(zip(
self.loader_node_name,
self.loader_var_use,
self.loader_api_call)):
blocks_node_name = [
blk.to(self.device) for blk in blocks_node_name
]
blocks_var_use = [
blk.to(self.device) for blk in blocks_var_use
]
blocks_api_call = [
blk.to(self.device) for blk in blocks_api_call
]
logits_node_name, labels_node_name = self._logits_node_name(
input_nodes_node_name, seeds_node_name, blocks_node_name)
logits_var_use, labels_var_use = self._logits_var_use(
input_nodes_var_use, seeds_var_use, blocks_var_use)
logits_api_call, labels_api_call = self._logits_api_call(
input_nodes_api_call, seeds_api_call, blocks_api_call)
train_acc_node_name = compute_accuracy(
logits_node_name.argmax(dim=1), labels_node_name)
train_acc_var_use = compute_accuracy(
logits_var_use.argmax(dim=1), labels_var_use)
train_acc_api_call = compute_accuracy(
logits_api_call.argmax(dim=1), labels_api_call)
train_logits = torch.cat(
[logits_node_name, logits_var_use, logits_api_call], 0)
train_labels = torch.cat(
[labels_node_name, labels_var_use, labels_api_call], 0)
logp = nn.functional.log_softmax(train_logits, 1)
loss = nn.functional.nll_loss(logp, train_labels)
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
self.write_summary(
{
"Loss": loss,
"Accuracy/train/node_name_vs_batch":
train_acc_node_name,
"Accuracy/train/var_use_vs_batch": train_acc_var_use,
"Accuracy/train/api_call_vs_batch": train_acc_api_call
}, self.batch)
self.batch += 1
self.eval()
with torch.set_grad_enabled(False):
_, val_acc_node_name, val_acc_var_use, val_acc_api_call = self._evaluate_objectives(
self.val_loader_node_name, self.val_loader_var_use,
self.val_loader_api_call, self.neg_sampling_factor)
_, test_acc_node_name, test_acc_var_use, test_acc_api_call = self._evaluate_objectives(
self.test_loader_node_name, self.test_loader_var_use,
self.test_loader_api_call, self.neg_sampling_factor)
self.train()
end = time()
self.best_score.track_best(epoch=epoch,
loss=loss.item(),
train_acc_node_name=train_acc_node_name,
val_acc_node_name=val_acc_node_name,
test_acc_node_name=test_acc_node_name,
train_acc_var_use=train_acc_var_use,
val_acc_var_use=val_acc_var_use,
test_acc_var_use=test_acc_var_use,
train_acc_api_call=train_acc_api_call,
val_acc_api_call=val_acc_api_call,
test_acc_api_call=test_acc_api_call,
time=end - start)
if self.do_save:
self.save_checkpoint(self.model_base_path)
self.write_summary(
{
"Accuracy/test/node_name_vs_batch": test_acc_node_name,
"Accuracy/test/var_use_vs_batch": test_acc_var_use,
"Accuracy/test/api_call_vs_batch": test_acc_api_call,
"Accuracy/val/node_name_vs_batch": val_acc_node_name,
"Accuracy/val/var_use_vs_batch": val_acc_var_use,
"Accuracy/val/api_call_vs_batch": val_acc_api_call
}, self.batch)
self.write_hyperparams(
{
"Loss/train_vs_epoch": loss,
"Accuracy/train/node_name_vs_epoch": train_acc_node_name,
"Accuracy/train/var_use_vs_epoch": train_acc_var_use,
"Accuracy/train/api_call_vs_epoch": train_acc_api_call,
"Accuracy/test/node_name_vs_epoch": test_acc_node_name,
"Accuracy/test/var_use_vs_epoch": test_acc_var_use,
"Accuracy/test/api_call_vs_epoch": test_acc_api_call,
"Accuracy/val/node_name_vs_epoch": val_acc_node_name,
"Accuracy/val/var_use_vs_epoch": val_acc_var_use,
"Accuracy/val/api_call_vs_epoch": val_acc_api_call
}, self.epoch)
self.lr_scheduler.step()
def save_checkpoint(self,
checkpoint_path=None,
checkpoint_name=None,
**kwargs):
checkpoint_path = join(checkpoint_path, "saved_state.pt")
param_dict = {
'graph_model': self.graph_model.state_dict(),
'node_embedder': self.node_embedder.state_dict(),
'ee_node_name': self.ee_node_name.state_dict(),
'ee_var_use': self.ee_var_use.state_dict(),
# 'ee_api_call': self.ee_api_call.state_dict(),
"lp_node_name": self.lp_node_name.state_dict(),
"lp_var_use": self.lp_var_use.state_dict(),
"lp_api_call": self.lp_api_call.state_dict(),
"epoch": self.epoch,
"batch": self.batch
}
if len(kwargs) > 0:
param_dict.update(kwargs)
torch.save(param_dict, checkpoint_path)
def restore_from_checkpoint(self, checkpoint_path):
checkpoint = torch.load(join(checkpoint_path, "saved_state.pt"))
self.graph_model.load_state_dict(checkpoint['graph_model'])
self.ee_node_name.load_state_dict(checkpoint['ee_node_name'])
self.ee_var_use.load_state_dict(checkpoint['ee_var_use'])
# self.ee_api_call.load_state_dict(checkpoint['ee_api_call'])
self.lp_node_name.load_state_dict(checkpoint['lp_node_name'])
self.lp_var_use.load_state_dict(checkpoint['lp_var_use'])
self.lp_api_call.load_state_dict(checkpoint['lp_api_call'])
self.epoch = checkpoint['epoch']
self.batch = checkpoint['batch']
logging.info(f"Restored from epoch {checkpoint['epoch']}")
def final_evaluation(self):
with torch.set_grad_enabled(False):
loss, train_acc_node_name, train_acc_var_use, train_acc_api_call = self._evaluate_objectives(
self.loader_node_name, self.loader_var_use,
self.loader_api_call, 1)
_, val_acc_node_name, val_acc_var_use, val_acc_api_call = self._evaluate_objectives(
self.val_loader_node_name, self.val_loader_var_use,
self.val_loader_api_call, 1)
_, test_acc_node_name, test_acc_var_use, test_acc_api_call = self._evaluate_objectives(
self.test_loader_node_name, self.test_loader_var_use,
self.test_loader_api_call, 1)
scores = {
# "loss": loss.item(),
"train_acc_node_name": train_acc_node_name,
"val_acc_node_name": val_acc_node_name,
"test_acc_node_name": test_acc_node_name,
"train_acc_var_use": train_acc_var_use,
"val_acc_var_use": val_acc_var_use,
"test_acc_var_use": test_acc_var_use,
"train_acc_api_call": train_acc_api_call,
"val_acc_api_call": val_acc_api_call,
"test_acc_api_call": test_acc_api_call,
}
print(
f'Final Eval : node name Train Acc {scores["train_acc_node_name"]:.4f}, '
f'node name Val Acc {scores["val_acc_node_name"]:.4f}, '
f'node name Test Acc {scores["test_acc_node_name"]:.4f}, '
f'var use Train Acc {scores["train_acc_var_use"]:.4f}, '
f'var use Val Acc {scores["val_acc_var_use"]:.4f}, '
f'var use Test Acc {scores["test_acc_var_use"]:.4f}, '
f'api call Train Acc {scores["train_acc_api_call"]:.4f}, '
f'api call Val Acc {scores["val_acc_api_call"]:.4f}, '
f'api call Test Acc {scores["test_acc_api_call"]:.4f}')
return scores
def eval(self):
self.graph_model.eval()
self.ee_node_name.eval()
self.ee_var_use.eval()
# self.ee_api_call.eval()
self.lp_node_name.eval()
self.lp_var_use.eval()
self.lp_api_call.eval()
def train(self):
self.graph_model.train()
self.ee_node_name.train()
self.ee_var_use.train()
# self.ee_api_call.eval()
self.lp_node_name.train()
self.lp_var_use.train()
self.lp_api_call.train()
def to(self, device):
self.graph_model.to(device)
self.ee_node_name.to(device)
self.ee_var_use.to(device)
# self.ee_api_call.to(device)
self.lp_node_name.to(device)
self.lp_var_use.to(device)
self.lp_api_call.to(device)
def get_embeddings(self):
# self.graph_model.g.nodes["function"].data.keys()
nodes = self.graph_model.g.nodes
node_embs = {
ntype: self.node_embedder(node_type=ntype,
node_ids=nodes[ntype].data['typed_id'],
train_embeddings=False)
for ntype in self.graph_model.g.ntypes
}
h = self.graph_model.inference(batch_size=256,
device='cpu',
num_workers=0,
x=node_embs)
original_id = []
global_id = []
embeddings = []
for ntype in self.graph_model.g.ntypes:
embeddings.append(h[ntype])
original_id.extend(nodes[ntype].data['original_id'].tolist())
global_id.extend(nodes[ntype].data['global_graph_id'].tolist())
embeddings = torch.cat(embeddings, dim=0).detach().numpy()
return [Embedder(dict(zip(original_id, global_id)), embeddings)]
def train(init_dict,
n_iterations,
transformer,
sim_data_node=None,
is_hsearch=False,
tb_writer=None,
print_every=10):
start = time.time()
predictor = init_dict['model']
data_loaders = init_dict['data_loaders']
optimizer = init_dict['optimizer']
metrics = init_dict['metrics']
lr_sch = ExponentialLR(optimizer, gamma=0.98)
if tb_writer:
tb_writer = tb_writer()
best_model_wts = predictor.state_dict()
best_score = -10000
best_epoch = -1
terminate_training = False
n_epochs = n_iterations // len(data_loaders['train'])
criterion = torch.nn.MSELoss()
# Since during hyperparameter search values that could cause CUDA memory exception could be sampled
# we want to ignore such values and find others that are workable within the memory constraints.
with contextlib.suppress(Exception if is_hsearch else DummyException):
for epoch in range(n_epochs):
eval_scores = []
for phase in ['train', 'val' if is_hsearch else 'test']:
if phase == 'train':
predictor.train()
else:
predictor.eval()
losses = []
metrics_dict = defaultdict(list)
for batch in tqdm(
data_loaders[phase],
desc=f'Phase: {phase}, epoch={epoch + 1}/{n_epochs}'
):
batch = np.array(batch)
x = batch[:, 0]
y_true = batch[:, 1]
with torch.set_grad_enabled(phase == 'train'):
y_true = torch.from_numpy(
y_true.reshape(-1, 1).astype(
np.float)).float().to(device)
y_pred = predictor(x)
loss = criterion(y_pred, y_true)
losses.append(loss.item())
# Perform evaluation using the given metrics
eval_dict = {}
score = ExpertTrainer.evaluate(
eval_dict,
transformer.inverse_transform(
y_true.cpu().detach().numpy()),
transformer.inverse_transform(
y_pred.cpu().detach().numpy()), metrics)
for m in eval_dict:
if m in metrics:
metrics_dict[m].append(eval_dict[m])
else:
metrics_dict[m] = [eval_dict[m]]
# Update weights
if phase == 'train':
optimizer.zero_grad()
loss.backward()
optimizer.step()
else:
eval_scores.append(score)
metrics_dict = {
k: np.mean(metrics_dict[k])
for k in metrics_dict
}
if epoch % print_every == 0:
print(
f'{phase}: epoch={epoch + 1}/{n_epochs}, loss={np.mean(losses)}, metrics={metrics_dict}'
)
if phase == 'train':
lr_sch.step()
# Checkpoint
score = np.mean(eval_scores)
if score > best_score:
best_score = score
best_model_wts = copy.deepcopy(predictor.state_dict())
best_epoch = epoch
predictor.load_state_dict(best_model_wts)
print(f'Time elapsed: {time_since(start)}')
return {'model': predictor, 'score': best_score, 'epoch': best_epoch}
output, c_D, _ = D(fake)
lossG = bce_loss(output, y_real)
lossinfo = mse_loss(c_D, c)
D_G_z2 = output.data.mean()
return lossG, lossinfo, D_G_z2
# start training
D.train()
for epoch in range(FG.num_epochs):
printers['lr']('D', epoch, optimizerD.param_groups[0]['lr'])
printers['lr']('G', epoch, optimizerG.param_groups[0]['lr'])
printers['lr']('info', epoch, optimizerinfo.param_groups[0]['lr'])
timer.tic()
#lr schedular
if (epoch+1)%100 == 0:
schedularD.step()
schedularG.step()
schedularinfo.step()
G.train()
for step, data in enumerate(trainloader):
D.zero_grad()
G.zero_grad()
x = data['image'].float().cuda(device, non_blocking=True)
y = data['target'].float().cuda(device, non_blocking=True)
batch_size = x.size(0)
# set veriables
z = torch.rand(batch_size, FG.z_dim).float().cuda(device, non_blocking=True)
c = torch.from_numpy(np.random.uniform(-1, 1, size=(batch_size,\
c_code))).float().cuda(device, non_blocking=True)
def train(data_dir, net_path=None, save_dir='pretrained'):
#从文件中读取图像数据集
seq_dataset = GOT10k(data_dir, subset='train', return_meta=False)
#定义图像预处理方法
transforms = SiamFCTransforms(
exemplar_sz=cfg.exemplar_sz, #127
instance_sz=cfg.instance_sz, #255
context=cfg.context) #0.5
#从读取的数据集每个视频序列配对训练图像并进行预处理，裁剪等
train_dataset = GOT10kDataset(seq_dataset, transforms)
#加载训练数据集
loader_dataset = DataLoader(
dataset=train_dataset,
batch_size=cfg.batch_size,
shuffle=True,
num_workers=cfg.num_workers,
pin_memory=True,
drop_last=True,
)
#初始化训练网络
cuda = torch.cuda.is_available() #支持GPU为True
device = torch.device('cuda:0' if cuda else 'cpu') #cuda设备号为0
model = AlexNet(init_weight=True)
corr = _corr()
model = model.to(device)
corr = corr.to(device)
# 设置损失函数和标签
logist_loss = BalancedLoss()
labels = _create_labels(
size=[cfg.batch_size, 1, cfg.response_sz - 2, cfg.response_sz - 2])
labels = torch.from_numpy(labels).to(device).float()
#建立优化器，设置指数变化的学习率
optimizer = optim.SGD(
model.parameters(),
lr=cfg.initial_lr, #初始化的学习率，后续会不断更新
weight_decay=cfg.weight_decay, #λ=5e-4，正则化
momentum=cfg.momentum) #v(now)=−dx∗lr+v(last)∗momemtum
gamma = np.power( #np.power(a,b) 返回a^b
cfg.ultimate_lr / cfg.initial_lr, 1.0 / cfg.epoch_num)
lr_scheduler = ExponentialLR(optimizer,
gamma) #指数形式衰减，lr=initial_lr*(gamma^epoch)
"""————————————————————————开始训练——————————————————————————"""
if not os.path.exists(save_dir):
os.makedirs(save_dir)
start_epoch = 1
#接着上一次训练，提取训练结束保存的net、optimizer、epoch参数
if net_path is not None:
checkpoint = torch.load(net_path)
if 'epoch' in checkpoint:
start_epoch = checkpoint['epoch'] + 1
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
else:
model.load_state_dict(checkpoint)
del checkpoint
torch.cuda.empty_cache() #缓存清零
print("loaded checkpoint!!!")
for epoch in range(start_epoch, cfg.epoch_num + 1):
model.train()
#遍历训练集
for it, batch in enumerate(tqdm(loader_dataset)):
z = batch[0].to(device,
non_blocking=cuda) # z.shape=([8,3,127,127])
x = batch[1].to(device,
non_blocking=cuda) # x.shape=([8,3,239,239])
#输入网络后通过损失函数
z, x = model(z), model(x)
responses = corr(
z, x
) * cfg.out_reduce # 返回的是heatmap的响应表15x15 因为x是239x239 [8,1,15,15]
loss = logist_loss(responses, labels)
#back propagation反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# if (it+1) % 20 ==0:
print('Epoch: {}[{}/{}] Loss: {:.5f} lr: {:.2e}'.format(
epoch, it + 1, len(loader_dataset), loss.item(),
optimizer.param_groups[0]['lr']))
#更新学习率 (每个epoch)
lr_scheduler.step()
#save checkpoint 做玩1个epoch后保存1个输出模型
if not os.path.exists(save_dir):
os.makedirs(save_dir)
save_path = os.path.join(save_dir, 'siamfc_alexnet_e%d.pth' % (epoch))
torch.save(
{
'epoch': epoch,
'model': model.state_dict(),
'optimizer': optimizer.state_dict()
}, save_path)
class TrackerSiamFC(Tracker):
def __init__(self, imagefile, region):
super(TrackerSiamFC, self).__init__(name='SiamFC',
is_deterministic=True)
self.cfg = self.parse_args()
# setup GPU device if available
self.cuda = torch.cuda.is_available()
self.device = torch.device('cuda:0' if self.cuda else 'cpu')
# setup model
self.net = SiamFC()
net_path = "/home/user/siamfc/pretrained/siamfc_new/model_e1_BEST.pth"
if net_path is not None:
self.net.load_state_dict(
torch.load(net_path,
map_location=lambda storage, loc: storage))
self.net = self.net.to(self.device)
# setup optimizer
self.optimizer = optim.SGD(self.net.parameters(),
lr=self.cfg.initial_lr,
weight_decay=self.cfg.weight_decay,
momentum=self.cfg.momentum)
# setup lr scheduler
self.lr_scheduler = ExponentialLR(self.optimizer,
gamma=self.cfg.lr_decay)
self.cf_influence = 0.11
bbox = convert_bbox_format(region, 'center-based')
bbox = (bbox.x, bbox.y, bbox.width, bbox.height)
image = Image.open(imagefile)
self.init(bbox, image)
def parse_args(self):
# default parameters
cfg = {
# inference parameters
'exemplar_sz': 127,
'instance_sz': 255,
'context': 0.5,
'scale_num': 3,
'scale_step': 1.0375,
'scale_lr': 0.59,
'scale_penalty': 0.9745,
'window_influence': 0.176, #change here 0.176 -> 0.1
'response_sz': 17,
'response_up': 16,
'total_stride': 8,
'adjust_scale': 0.001,
# train parameters
'initial_lr': 0.001, #change here 0.01->0.001
'lr_decay': 0.8685113737513527,
'weight_decay': 5e-4,
'momentum': 0.9,
'r_pos': 16,
'r_neg': 0
}
return namedtuple('GenericDict', cfg.keys())(**cfg)
def init(self, image, box):
self.net.is_train = False
image = np.asarray(image)
# convert box to 0-indexed and center based [y, x, h, w]
box = np.array([
box[1] - 1 + (box[3] - 1) / 2, box[0] - 1 +
(box[2] - 1) / 2, box[3], box[2]
],
dtype=np.float32)
self.center, self.target_sz = box[:2], box[2:]
# create hanning window
self.upscale_sz = self.cfg.response_up * self.cfg.response_sz
self.hann_window = np.outer(np.hanning(self.upscale_sz),
np.hanning(self.upscale_sz))
self.hann_window /= self.hann_window.sum()
# search scale factors
self.scale_factors = self.cfg.scale_step**np.linspace(
-(self.cfg.scale_num // 2), self.cfg.scale_num // 2,
self.cfg.scale_num)
# exemplar and search sizes
context = self.cfg.context * np.sum(self.target_sz)
self.z_sz = np.sqrt(np.prod(self.target_sz + context))
self.x_sz = self.z_sz * \
self.cfg.instance_sz / self.cfg.exemplar_sz
# exemplar image
self.avg_color = np.mean(image, axis=(0, 1))
#TODO: change here
exemplar_image = self._crop_and_resize(image,
self.center,
self.z_sz,
out_size=self.cfg.exemplar_sz,
pad_color=self.avg_color)
exemplar_image_cf = self._crop_and_resize(
image,
self.center,
self.x_sz,
out_size=self.cfg.instance_sz,
pad_color=self.avg_color)
# exemplar features
exemplar_image = Image_to_Tensor(exemplar_image).to(
self.device).unsqueeze(0)
exemplar_image_cf = Image_to_Tensor(exemplar_image_cf).to(
self.device).unsqueeze(0)
#exemplar_image = torch.from_numpy(exemplar_image).to(
#self.device).permute([2, 0, 1]).unsqueeze(0).float()
with torch.set_grad_enabled(False):
self.net.eval()
_, self.kernel = self.net.features(exemplar_image)
self.kernel = self.kernel.repeat(3, 1, 1, 1)
self.net.update(exemplar_image_cf)
def track(self, imagefile):
image = Image.open(imagefile)
self.update(image)
bbox = Rectangle(self.center[0], self.center[1], self.target_sz[0],
self.target_sz[1])
bbox = convert_bbox_format(bbox, 'top-left-based')
return bbox
def update(self, image):
self.net.is_train = False
image = np.asarray(image)
# search images
instance_images = [
self._crop_and_resize(image,
self.center,
self.x_sz * f,
out_size=self.cfg.instance_sz,
pad_color=self.avg_color)
for f in self.scale_factors
]
instance_images = [
Image_to_Tensor(f).to(self.device).unsqueeze(0).squeeze(0)
for f in instance_images
]
instance_images = torch.stack(instance_images)
# responses
with torch.set_grad_enabled(False):
self.net.eval()
#TODO: change here
#_, instances = self.net.features(instance_images)
#responses = F.conv2d(instances, self.kernel) * 0.001
responses, cf_responses = self.net.get_response(
self.kernel, instance_images)
responses = responses.squeeze(1).cpu().numpy()
cf_responses = cf_responses.squeeze(1).cpu().numpy()
#print(np.unravel_index(cf_responses[1].argmax(), cf_responses[1].shape))
cf_responses = np.roll(cf_responses,
int(np.floor(float(251) / 2.) - 1),
axis=1)
cf_responses = np.roll(cf_responses,
int(np.floor(float(251) / 2.) - 1),
axis=2)
#print(np.unravel_index(cf_responses[1].argmax(), cf_responses[1].shape))
#cv2.imshow("tset", cf_responses[1])
#cv2.waitKey(1000)
# upsample responses and penalize scale changes
#cf-----------------------------------------------------------------------
cf_responses = np.stack([
cv2.resize(t, (510, 510), interpolation=cv2.INTER_CUBIC)
for t in cf_responses
],
axis=0)
#cf_responses[:self.cfg.scale_num // 2] *= self.cfg.scale_penalty
#cf_responses[self.cfg.scale_num // 2 + 1:] *= self.cfg.scale_penalty
#cf_scale_id = np.argmax(np.amax(cf_responses, axis=(1, 2)))
#cf_response = cf_responses[cf_scale_id]
#cf_loc = np.unravel_index(cf_response.argmax(), cf_response.shape)
#print(cf_loc)
#cf_disp_in_response = np.array(cf_loc) - 255 // 2
#cf_disp_in_image = cf_disp_in_response * self.x_sz * \
#self.scale_factors[cf_scale_id] / self.cfg.instance_sz
#print(cf_disp_in_image)
cf_responses = cf_responses[:, 119:391, 119:391]
#cv2.imshow("tset", cf_responses[1])
#cv2.waitKey(1000)
#-------------------------------------------------------------------------
#siamfc-------------------------------------------------------------------
responses = np.stack([
cv2.resize(t, (self.upscale_sz, self.upscale_sz),
interpolation=cv2.INTER_CUBIC) for t in responses
],
axis=0)
responses[:self.cfg.scale_num // 2] *= self.cfg.scale_penalty
responses[self.cfg.scale_num // 2 + 1:] *= self.cfg.scale_penalty
# peak scale
scale_id = np.argmax(np.amax(responses, axis=(1, 2)))
# peak location
response = responses[scale_id]
cf_response = cf_responses[scale_id]
cf_response -= cf_response.min()
cf_response /= cf_response.sum()
response -= response.min()
response /= response.sum() + 1e-16
#response = (1 - self.cfg.window_influence) * response + \
#self.cfg.window_influence * self.hann_window
response = (1 - self.cfg.window_influence) * response + \
self.cfg.window_influence * self.hann_window + self.cf_influence * cf_response
loc = np.unravel_index(response.argmax(), response.shape)
# locate target center
disp_in_response = np.array(loc) - self.upscale_sz // 2
disp_in_instance = disp_in_response * \
self.cfg.total_stride / self.cfg.response_up
disp_in_image = disp_in_instance * self.x_sz * \
self.scale_factors[scale_id] / self.cfg.instance_sz
#--------------------------------------------------------------------------
self.center += disp_in_image
#self.center += cf_disp_in_image
# update target size
scale = (1 - self.cfg.scale_lr) * 1.0 + \
self.cfg.scale_lr * self.scale_factors[scale_id]
self.target_sz *= scale
self.z_sz *= scale
self.x_sz *= scale
# update cf
exemplar_image_cf = self._crop_and_resize(
image,
self.center,
self.x_sz,
out_size=self.cfg.instance_sz,
pad_color=self.avg_color)
exemplar_image_cf = Image_to_Tensor(exemplar_image_cf).to(
self.device).unsqueeze(0)
self.net.update(exemplar_image_cf, lr=0.01)
# return 1-indexed and left-top based bounding box
box = np.array([
self.center[1] + 1 - (self.target_sz[1] - 1) / 2,
self.center[0] + 1 - (self.target_sz[0] - 1) / 2,
self.target_sz[1], self.target_sz[0]
])
return box
def step(self, batch, backward=True, update_lr=False):
self.net.is_train = True
if backward:
self.net.train()
if update_lr:
self.lr_scheduler.step()
else:
self.net.eval()
z = batch[0].to(self.device)
x = batch[1].to(self.device)
label = batch[2].to(self.device)
with torch.set_grad_enabled(backward):
responses, out2 = self.net(z, x)
labels, weights = self._create_labels(responses.size())
loss1 = F.binary_cross_entropy_with_logits(responses,
labels,
weight=weights,
size_average=True)
loss2 = F.mse_loss(out2, label, size_average=True)
loss = loss1 + loss2 * 5
if backward:
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
def _crop_and_resize(self, image, center, size, out_size, pad_color):
# convert box to corners (0-indexed)
size = round(size)
corners = np.concatenate((np.round(center - (size - 1) / 2),
np.round(center - (size - 1) / 2) + size))
corners = np.round(corners).astype(int)
# pad image if necessary
pads = np.concatenate((-corners[:2], corners[2:] - image.shape[:2]))
npad = max(0, int(pads.max()))
if npad > 0:
image = cv2.copyMakeBorder(image,
npad,
npad,
npad,
npad,
cv2.BORDER_CONSTANT,
value=pad_color)
# crop image patch
corners = (corners + npad).astype(int)
patch = image[corners[0]:corners[2], corners[1]:corners[3]]
# resize to out_size
patch = cv2.resize(patch, (out_size, out_size))
return patch
def _create_labels(self, size):
# skip if same sized labels already created
if hasattr(self, 'labels') and self.labels.size() == size:
return self.labels, self.weights
def logistic_labels(x, y, r_pos, r_neg):
dist = np.abs(x) + np.abs(y) # block distance
labels = np.where(
dist <= r_pos, np.ones_like(x),
np.where(dist < r_neg,
np.ones_like(x) * 0.5, np.zeros_like(x)))
return labels
# distances along x- and y-axis
n, c, h, w = size
x = np.arange(w) - w // 2
y = np.arange(h) - h // 2
x, y = np.meshgrid(x, y)
# create logistic labels
r_pos = self.cfg.r_pos / self.cfg.total_stride
r_neg = self.cfg.r_neg / self.cfg.total_stride
labels = logistic_labels(x, y, r_pos, r_neg)
# pos/neg weights
pos_num = np.sum(labels == 1)
neg_num = np.sum(labels == 0)
weights = np.zeros_like(labels)
weights[labels == 1] = 0.5 / pos_num
weights[labels == 0] = 0.5 / neg_num
weights *= pos_num + neg_num
# repeat to size
labels = labels.reshape((1, 1, h, w))
weights = weights.reshape((1, 1, h, w))
labels = np.tile(labels, (n, c, 1, 1))
weights = np.tile(weights, [n, c, 1, 1])
# convert to tensors
self.labels = torch.from_numpy(labels).to(self.device).float()
self.weights = torch.from_numpy(weights).to(self.device).float()
return self.labels, self.weights
class Experiment:
def __init__(self,
*,
model,
learning_rate=0.0005,
embedding_dim,
num_iterations,
batch_size=128,
decay_rate=0.,
conv_out=2,
projection_size=10,
input_dropout=0.4,
feature_map_dropout=0.4,
hidden_dropout=0.4,
label_smoothing=0.1,
cuda=True):
self.model = model
self.learning_rate = learning_rate
self.num_iterations = num_iterations
self.batch_size = batch_size
self.decay_rate = decay_rate
self.label_smoothing = label_smoothing
self.cuda = torch.cuda.is_available()
self.entity_idxs, self.relation_idxs, self.scheduler = None, None, None
# Params stored in kwargs for creating pretrained models with unique names.
self.kwargs = {
'embedding_dim': embedding_dim,
'learning_rate': learning_rate,
'batch_size': batch_size,
'conv_out': conv_out,
'input_dropout': input_dropout,
'hidden_dropout': hidden_dropout,
'projection_size': projection_size,
'feature_map_dropout': feature_map_dropout,
'label_smoothing': label_smoothing,
'decay_rate': decay_rate
}
self.storage_path, _ = create_experiment_folder()
self.logger = create_logger(name=self.model, p=self.storage_path)
def get_data_idxs(self, data):
data_idxs = [
(self.entity_idxs[data[i][0]], self.relation_idxs[data[i][1]],
self.entity_idxs[data[i][2]]) for i in range(len(data))
]
return data_idxs
def get_er_vocab(self, data):
er_vocab = defaultdict(list)
for triple in data:
er_vocab[(triple[0], triple[1])].append(triple[2])
return er_vocab
def get_batch(self, er_vocab, er_vocab_pairs, idx):
batch = er_vocab_pairs[idx:idx + self.batch_size]
targets = np.zeros((len(batch), len(d.entities)))
for idx, pair in enumerate(batch):
targets[idx, er_vocab[pair]] = 1.
targets = torch.FloatTensor(targets)
if self.cuda:
targets = targets.cuda()
return np.array(batch), targets
def evaluate(self, model, data):
hits = []
ranks = []
for i in range(10):
hits.append([])
test_data_idxs = self.get_data_idxs(data)
er_vocab = self.get_er_vocab(self.get_data_idxs(d.data))
for i in range(0, len(test_data_idxs), self.batch_size):
data_batch, _ = self.get_batch(er_vocab, test_data_idxs, i)
e1_idx = torch.tensor(data_batch[:, 0])
r_idx = torch.tensor(data_batch[:, 1])
e2_idx = torch.tensor(data_batch[:, 2])
if self.cuda:
e1_idx = e1_idx.cuda()
r_idx = r_idx.cuda()
e2_idx = e2_idx.cuda()
predictions = model.forward(e1_idx, r_idx)
for j in range(data_batch.shape[0]):
filt = er_vocab[(data_batch[j][0], data_batch[j][1])]
target_value = predictions[j, e2_idx[j]].item()
predictions[j, filt] = 0.0
predictions[j, e2_idx[j]] = target_value
sort_values, sort_idxs = torch.sort(predictions,
dim=1,
descending=True)
sort_idxs = sort_idxs.cpu().numpy()
for j in range(data_batch.shape[0]):
rank = np.where(sort_idxs[j] == e2_idx[j].item())[0][0]
ranks.append(rank + 1)
for hits_level in range(10):
if rank <= hits_level:
hits[hits_level].append(1.0)
else:
hits[hits_level].append(0.0)
self.logger.info('Hits @10: {0}'.format(np.mean(hits[9])))
self.logger.info('Hits @3: {0}'.format(np.mean(hits[2])))
self.logger.info('Hits @1: {0}'.format(np.mean(hits[0])))
# print('Mean rank: {0}'.format(np.mean(ranks)))
self.logger.info('Mean reciprocal rank: {0}'.format(
np.mean(1. / np.array(ranks))))
def train_and_eval(self, d_info):
self.entity_idxs = {d.entities[i]: i for i in range(len(d.entities))}
self.relation_idxs = {
d.relations[i]: i
for i in range(len(d.relations))
}
train_data_idxs = self.get_data_idxs(d.train_data)
self.kwargs.update({
'num_entities': len(self.entity_idxs),
'num_relations': len(self.relation_idxs)
})
self.kwargs.update(d_info)
self.logger.info("Info pertaining to dataset:{0}".format(
d_info['dataset']))
self.logger.info("Number of triples in training data:{0}".format(
len(d.train_data)))
self.logger.info("Number of triples in validation data:{0}".format(
len(d.valid_data)))
self.logger.info("Number of triples in testing data:{0}".format(
len(d.test_data)))
self.logger.info("Number of entities:{0}".format(len(
self.entity_idxs)))
self.logger.info("Number of relations:{0}".format(
len(self.relation_idxs)))
self.logger.info("HyperParameter Settings:{0}".format(self.kwargs))
model = None
if self.model == 'Conex':
model = ConEx(self.kwargs)
elif self.model == 'Distmult':
model = DistMult(self.kwargs)
elif self.model == 'Complex':
model = Complex(self.kwargs)
elif self.model == 'Conve':
model = ConvE(self.kwargs)
elif self.model == 'Tucker':
model = TuckER(self.kwargs)
elif self.model == 'HypER':
model = HypER(self.kwargs)
else:
raise ValueError
self.logger.info(model)
if self.cuda:
model.cuda()
model.init()
opt = torch.optim.Adam(model.parameters(), lr=self.learning_rate)
if self.decay_rate:
self.scheduler = ExponentialLR(opt, self.decay_rate)
er_vocab = self.get_er_vocab(train_data_idxs)
er_vocab_pairs = list(er_vocab.keys())
self.logger.info("{0} starts training".format(model.name))
num_param = sum([p.numel() for p in model.parameters()])
self.logger.info("'Number of free parameters: {0}".format(num_param))
for it in range(1, self.num_iterations + 1):
model.train()
losses = []
np.random.shuffle(er_vocab_pairs)
for j in range(0, len(er_vocab_pairs), self.batch_size):
data_batch, targets = self.get_batch(er_vocab, er_vocab_pairs,
j)
opt.zero_grad()
e1_idx = torch.tensor(data_batch[:, 0])
r_idx = torch.tensor(data_batch[:, 1])
if self.cuda:
e1_idx = e1_idx.cuda()
r_idx = r_idx.cuda()
predictions = model.forward(e1_idx, r_idx)
if self.label_smoothing:
targets = ((1.0 - self.label_smoothing) *
targets) + (1.0 / targets.size(1))
loss = model.loss(predictions, targets)
loss.backward()
opt.step()
losses.append(loss.item())
if self.decay_rate:
self.scheduler.step()
if it % 500 == 0:
self.logger.info('Iteration:{0} with Average loss{1}'.format(
it, np.mean(losses)))
model.eval(
) # Turns evaluation mode on, i.e., dropouts are turned off.
with torch.no_grad(): # Important:
self.logger.info("Validation:")
self.evaluate(model, d.valid_data)
with open(self.storage_path + '/settings.json',
'w') as file_descriptor:
json.dump(self.kwargs, file_descriptor)
self.logger.info("Testing:")
self.evaluate(model, d.test_data)
torch.save(model.state_dict(), self.storage_path + '/model.pt')
'reconstructions': wandb.Image(recons, caption = 'reconstructions'),
'hard reconstructions': wandb.Image(hard_recons, caption = 'hard reconstructions'),
'codebook_indices': wandb.Histogram(codes),
'temperature': temp
}
save_model(f'./vae.pt')
wandb.save('./vae.pt')
# temperature anneal
temp = max(temp * math.exp(-ANNEAL_RATE * global_step), TEMP_MIN)
# lr decay
sched.step()
if i % 10 == 0:
lr = sched.get_last_lr()[0]
print(epoch, i, f'lr - {lr:6f} loss - {loss.item()}')
logs = {
**logs,
'epoch': epoch,
'iter': i,
'loss': loss.item(),
'lr': lr
}
wandb.log(logs)
global_step += 1
def main():
global args, best_loss
global writer, csv_writer
global device, kwargs
base_model = load_model(arch=args.arch, pretrained=True)
if args.use_parallel:
model = FineTuneModelPool(base_model, args.arch, args.num_classes,
str(args.classifier_config))
model = torch.nn.DataParallel(model).to(device)
else:
model = FineTuneModelPool(base_model, args.arch, args.num_classes,
str(args.classifier_config)).to(device)
if args.use_parallel:
params = model.module.parameters()
mean = model.module.mean
std = model.module.std
else:
params = model.parameters()
mean = model.mean
std = model.std
if args.optimizer.startswith('adam'):
optimizer = torch.optim.Adam(
filter(lambda p: p.requires_grad, params),
# Only finetunable params
lr=args.lr)
elif args.optimizer.startswith('rmsprop'):
optimizer = torch.optim.RMSprop(
filter(lambda p: p.requires_grad, params),
# Only finetunable params
lr=args.lr)
elif args.optimizer.startswith('sgd'):
optimizer = torch.optim.SGD(
filter(lambda p: p.requires_grad, params),
# Only finetunable params
lr=args.lr)
else:
raise ValueError('Optimizer not supported')
# optionally resume from a checkpoint
loaded_from_checkpoint = False
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
if args.use_parallel:
model.load_state_dict(checkpoint['state_dict'])
else:
model.load_state_dict(checkpoint['state_dict'])
# optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint (epoch {})".format(
checkpoint['epoch']))
loaded_from_checkpoint = True
else:
print("=> no checkpoint found at '{}'".format(args.resume))
if args.predict:
pass
elif args.evaluate:
pass
else:
base_dset_kwargs = {
'mode': 'train',
'random_state': args.seed,
'fold': args.fold,
'size_ratio': args.size_ratio,
'preprocessing_type': args.preprocessing_type,
'fixed_size': (224, 224),
'prob': 0.2,
'mean': mean,
'std': std,
}
train_dataset, val_dataset, train_sampler, val_sampler, train_loader, val_loader = get_datasets(
base_dset_kwargs, args.batch_size)
if args.multi_class:
criterion = MultiClassBCELoss().to(device)
else:
criterion = nn.CrossEntropyLoss().cuda().to(device)
hard_dice_05 = HardDice(threshold=0.5)
if args.lr_regime == 'auto_decay':
print('Auto-decay')
scheduler = ExponentialLR(optimizer=optimizer,
gamma=0.99,
last_epoch=-1)
elif args.lr_regime == 'plateau_decay':
print('Plateau decay')
scheduler = ReduceLROnPlateau(optimizer=optimizer,
mode='min',
factor=0.5,
patience=5,
verbose=True)
elif args.lr_regime == 'clr':
print('CLR decay')
scheduler = CyclicLR(optimizer=optimizer,
base_lr=1e-4,
max_lr=1e-2,
step_size=1200,
mode='exp_range',
gamma=0.95)
for epoch in range(args.start_epoch, args.epochs):
# grow only if param specified
if args.epochs_grow_size > 0:
# grow each number of epochs
if (epoch + 1) % args.epochs_grow_size == 0:
# do not grow over
if train_dataset.size_ratio < 1.0:
# increase the current size ration by a factor of 2
# also decrease the batch size by 4
new_batch_size = int(train_loader.batch_size // 4)
train_dataset, val_dataset, train_sampler, val_sampler, train_loader, val_loader = get_datasets(
{
**base_dset_kwargs,
**{
'size_ratio': train_dataset.size_ratio * 2
}
}, new_batch_size)
# train for one epoch
train_loss, train_hard_dice_05, train_f1, train_acc1, train_acc5 = train(
train_loader, model, criterion, hard_dice_05, optimizer, epoch,
scheduler)
# evaluate on validation set
val_loss, val_hard_dice_05, val_f1, val_acc1, val_acc5 = validate(
val_loader, model, criterion, hard_dice_05)
if args.lr_regime == 'auto_decay':
scheduler.step()
elif args.lr_regime == 'plateau_decay':
scheduler.step(val_loss)
#============ TensorBoard logging ============#
# Log the scalar values
if args.tensorboard:
writer.add_scalars('epoch/epoch_losses', {
'train_loss': train_loss,
'val_loss': val_loss
}, epoch + 1)
if train_hard_dice_05 and val_hard_dice_05:
writer.add_scalars(
'epoch/epoch_hdice05', {
'train_hdice': train_hard_dice_05,
'val_hdice': val_hard_dice_05
}, epoch + 1)
if train_acc1 and val_acc1:
writer.add_scalars('epoch/epoch_acc1', {
'train_acc1': train_acc1,
'val_acc1': val_acc1
}, epoch + 1)
else:
train_acc1 = 0
val_acc1 = 0
if train_acc5 and val_acc5:
writer.add_scalars('epoch/epoch_acc5', {
'train_acc5': train_acc5,
'val_acc5': val_acc5
}, epoch + 1)
else:
train_acc5 = 0
val_acc5 = 0
if train_f1 and val_f1:
writer.add_scalars('epoch/epoch_f1', {
'train_f1': train_f1,
'val_f1': val_f1
}, epoch + 1)
else:
train_f1 = 0
val_f1 = 0
csv_writer.write({
'epoch': epoch + 1,
'train_acc1': train_acc1,
'val_acc1': val_acc1,
'train_acc5': train_acc5,
'val_acc5': val_acc5,
'train_f1': train_f1,
'val_f1': val_f1
})
# remember best prec@1 and save checkpoint
is_best = val_loss < best_loss
best_loss = min(val_loss, best_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'optimizer': optimizer.state_dict(),
'state_dict': model.state_dict(),
'best_loss': best_loss,
}, is_best,
'weights/{}_checkpoint.pth.tar'.format(str(args.lognumber)),
'weights/{}_best.pth.tar'.format(str(args.lognumber)))
def train_and_eval(self):
logger.info(
f'Training the {model_name} model with {dataset} knowledge graph ...'
)
self.entity_idxs = {d.entities[i]: i for i in range(len(d.entities))}
self.relation_idxs = {
d.relations[i]: i
for i in range(len(d.relations))
}
train_triple_idxs = self.get_data_idxs(d.train_data)
train_triple_size = len(train_triple_idxs)
logger.info(f'Number of training data points: {train_triple_size}')
if model_name.lower() == "hype":
model = HypE(d, self.ent_vec_dim, self.rel_vec_dim, **self.kwargs)
elif model_name.lower() == "hyper":
model = HypER(d, self.ent_vec_dim, self.rel_vec_dim, **self.kwargs)
elif model_name.lower() == "distmult":
model = DistMult(d, self.ent_vec_dim, self.rel_vec_dim,
**self.kwargs)
elif model_name.lower() == "conve":
model = ConvE(d, self.ent_vec_dim, self.rel_vec_dim, **self.kwargs)
elif model_name.lower() == "complex":
model = ComplEx(d, self.ent_vec_dim, self.rel_vec_dim,
**self.kwargs)
logger.debug('model parameters: {}'.format(
{name: value.numel()
for name, value in model.named_parameters()}))
if self.cuda:
model.cuda()
model.init()
opt = torch.optim.Adam(model.parameters(), lr=self.learning_rate)
if self.decay_rate:
scheduler = ExponentialLR(opt, self.decay_rate)
er_vocab = self.get_er_vocab(train_triple_idxs)
er_vocab_pairs = list(er_vocab.keys())
er_vocab_pairs_size = len(er_vocab_pairs)
logger.info(
f'Number of entity-relational pairs: {er_vocab_pairs_size}')
logger.info('Starting Training ...')
for epoch in range(1, self.epochs + 1):
logger.info(f'Epoch: {epoch}')
model.train()
costs = []
np.random.shuffle(er_vocab_pairs)
for j in range(0, er_vocab_pairs_size, self.batch_size):
if j % (128 * 100) == 0:
logger.info(f'Batch: {j + 1} ...')
triples, targets = self.get_batch(er_vocab, er_vocab_pairs,
er_vocab_pairs_size, j)
opt.zero_grad()
e1_idx = torch.tensor(triples[:, 0])
r_idx = torch.tensor(triples[:, 1])
if self.cuda:
e1_idx = e1_idx.cuda()
r_idx = r_idx.cuda()
predictions = model.forward(e1_idx, r_idx)
if self.label_smoothing:
targets = ((1.0 - self.label_smoothing) *
targets) + (1.0 / targets.size(1))
cost = model.loss(predictions, targets)
cost.backward()
opt.step()
costs.append(cost.item())
if self.decay_rate:
scheduler.step()
logger.info(f'Mean training cost: {np.mean(costs)}')
if epoch % 10 == 0:
model.eval()
with torch.no_grad():
train_data = np.array(d.train_data)
train_data_map = {
'WN18': 10000,
'FB15k': 100000,
'WN18RR': 6068,
'FB15k-237': 35070
}
train_data_sample_size = train_data_map[dataset]
train_data = train_data[
np.random.choice(train_data.shape[0],
train_data_sample_size,
replace=False), :]
logger.info(f'Starting Evaluation: Training ...')
self.evaluate(model, train_data, epoch, 'training')
logger.info(f'Evaluation: Training complete!')
logger.info(f'Starting Evaluation: Validation ...')
self.evaluate(model, d.valid_data, epoch, 'validation')
logger.info(f'Evaluation: Validation complete!')
logger.info(f'Starting Evaluation: Test ...')
self.evaluate(model, d.test_data, epoch, 'testing')
logger.info(f'Evaluation: Test complete!')
logger.info('Checkpointing model ...')
torch.save(model.state_dict(), 'HypER.mc')
logger.info('Model checkpoint complete!')
logger.info('Saving final model ...')
torch.save(model.state_dict(), 'HypER.pt')
logger.info('Saving final model complete!')
def train(data_path,
entity_path,
relation_path,
entity_dict,
relation_dict,
neg_batch_size,
batch_size,
shuffle,
num_workers,
nb_epochs,
embedding_dim,
hidden_dim,
relation_dim,
gpu,
use_cuda,
patience,
freeze,
validate_every,
num_hops,
lr,
entdrop,
reldrop,
scoredrop,
l3_reg,
model_name,
decay,
ls,
w_matrix,
bn_list,
valid_data_path=None):
entities = np.load(entity_path)
relations = np.load(relation_path)
e, r = preprocess_entities_relations(entity_dict, relation_dict, entities,
relations)
entity2idx, idx2entity, embedding_matrix = prepare_embeddings(e)
data = process_text_file(data_path, split=False)
# data = pickle.load(open(data_path, 'rb'))
word2ix, idx2word, max_len = get_vocab(data)
hops = str(num_hops)
# print(idx2word)
# aditay
# print(idx2word.keys())
device = torch.device(gpu if use_cuda else "cpu")
dataset = DatasetMetaQA(data=data,
word2ix=word2ix,
relations=r,
entities=e,
entity2idx=entity2idx)
data_loader = DataLoaderMetaQA(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=num_workers)
model = RelationExtractor(embedding_dim=embedding_dim,
hidden_dim=hidden_dim,
vocab_size=len(word2ix),
num_entities=len(idx2entity),
relation_dim=relation_dim,
pretrained_embeddings=embedding_matrix,
freeze=freeze,
device=device,
entdrop=entdrop,
reldrop=reldrop,
scoredrop=scoredrop,
l3_reg=l3_reg,
model=model_name,
ls=ls,
w_matrix=w_matrix,
bn_list=bn_list)
model.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
scheduler = ExponentialLR(optimizer, decay)
optimizer.zero_grad()
best_score = -float("inf")
best_model = model.state_dict()
no_update = 0
for epoch in range(nb_epochs):
phases = []
for i in range(validate_every):
phases.append('train')
phases.append('valid')
for phase in phases:
if phase == 'train':
model.train()
if freeze == True:
# print('Freezing batch norm layers')
model.apply(set_bn_eval)
loader = tqdm(data_loader,
total=len(data_loader),
unit="batches")
running_loss = 0
for i_batch, a in enumerate(loader):
model.zero_grad()
question = a[0].to(device)
sent_len = a[1].to(device)
positive_head = a[2].to(device)
positive_tail = a[3].to(device)
loss = model(sentence=question,
p_head=positive_head,
p_tail=positive_tail,
question_len=sent_len)
loss.backward()
optimizer.step()
running_loss += loss.item()
loader.set_postfix(Loss=running_loss /
((i_batch + 1) * batch_size),
Epoch=epoch)
loader.set_description('{}/{}'.format(epoch, nb_epochs))
loader.update()
scheduler.step()
elif phase == 'valid':
model.eval()
eps = 0.0001
answers, score = validate(model=model,
data_path=valid_data_path,
word2idx=word2ix,
entity2idx=entity2idx,
device=device,
model_name=model_name)
if score > best_score + eps:
best_score = score
no_update = 0
best_model = model.state_dict()
print(
hops +
" hop Validation accuracy increased from previous epoch",
score)
_, test_score = validate(model=model,
data_path=test_data_path,
word2idx=word2ix,
entity2idx=entity2idx,
device=device,
model_name=model_name)
print('Test score for best valid so far:', test_score)
# writeToFile(answers, 'results_' + model_name + '_' + hops + '.txt')
suffix = ''
if freeze == True:
suffix = '_frozen'
checkpoint_path = '../../checkpoints/MetaQA/'
checkpoint_file_name = checkpoint_path + model_name + '_' + num_hops + suffix + ".pt"
print('Saving checkpoint to ', checkpoint_file_name)
torch.save(model.state_dict(), checkpoint_file_name)
elif (score < best_score + eps) and (no_update < patience):
no_update += 1
print(
"Validation accuracy decreases to %f from %f, %d more epoch to check"
% (score, best_score, patience - no_update))
elif no_update == patience:
print(
"Model has exceed patience. Saving best model and exiting"
)
torch.save(best_model,
checkpoint_path + "best_score_model.pt")
exit()
if epoch == nb_epochs - 1:
print(
"Final Epoch has reached. Stopping and saving model.")
torch.save(best_model,
checkpoint_path + "best_score_model.pt")
exit()
optimizer.zero_grad()
# torch to variable
X_train_variable = Variable(torch.from_numpy(X_train).float())
y_train_variable = Variable(torch.from_numpy(y_train).float())
X_val_variable = Variable(torch.from_numpy(X_val).float())
y_val_variable = Variable(torch.from_numpy(y_val).float())
# train model
mlp.train()
outputs = mlp(X_train_variable)
loss = criterion(outputs, y_train_variable)
loss.backward()
optimizer.step()
scheduler.step()
if epoch % 100 == 0:
print('Epoch [%d/%d], Validation Loss: %.4f' %
(epoch + 1, num_epochs, loss.data[0]))
# validate model
mlp.eval() # for the case of dropout-layer and Batch-normalization
outputs = mlp(X_val_variable)
validation_loss = criterion(outputs, y_val_variable)
#print(validation_loss.data[0])
#total_loss += validation_loss.data[0]
#print(total_loss)
print("End of one cross validation subset")
def main():
learning_rate = 0.0005
BCEloss = nn.BCELoss().cuda()
MSEloss = nn.MSELoss().cuda()
generator = model.Generator()
#generator.load_state_dict(t.load('checkpoints/4fx/epoch0_g.tar'))
discriminator = model.Discriminator()
#discriminator.load_state_dict(t.load('checkpoints/4fx/epoch0_d.tar'))
classified = model.Classified()
#classified.load_state_dict(t.load('checkpoints/4fx/epoch0_c.tar'))
generator = generator.cuda()
discriminator = discriminator.cuda()
classified = classified.cuda()
generator.train()
discriminator.train()
classified.train()
t1 = lo.NgsimDataset('data/5feature/TrainSet.mat')
t2 = lo.NgsimDataset('data/5feature/ValSet.mat')
trainDataloader = DataLoader(t1,
batch_size=128,
shuffle=True,
num_workers=8,
collate_fn=t1.collate_fn) #46272batch
valDataloader = DataLoader(t2,
batch_size=128,
shuffle=True,
num_workers=8,
collate_fn=t2.collate_fn) #6716batch
optimizer_g = optim.Adam(generator.parameters(), lr=learning_rate)
optimizer_d = optim.Adam(discriminator.parameters(), lr=learning_rate)
optimizer_c = optim.Adam(classified.parameters(), lr=learning_rate)
scheduler_g = ExponentialLR(optimizer_g, gamma=0.6)
scheduler_d = ExponentialLR(optimizer_d, gamma=0.5)
scheduler_c = ExponentialLR(optimizer_c, gamma=0.6)
file = open('./checkpoints/6f/loss.txt', 'w')
for epoch in range(6):
print("epoch:", epoch, 'lr', optimizer_d.param_groups[0]['lr'])
loss_gi1 = 0
loss_gix = 0
loss_gi3 = 0
loss_gi4 = 0
for idx, data in enumerate(trainDataloader):
hist, nbrs, mask, lat_enc, lon_enc, fut, op_mask, va, nbrsva, lane, nbrslane, dis, nbrsdis = data
hist = hist.cuda()
nbrs = nbrs.cuda()
mask = mask.cuda()
lat_enc = lat_enc.cuda()
lon_enc = lon_enc.cuda()
fut = fut.cuda()
op_mask = op_mask.cuda()
va = va.cuda()
nbrsva = nbrsva.cuda()
lane = lane.cuda()
nbrslane = nbrslane.cuda()
dis = dis.cuda()
nbrsdis = nbrsdis.cuda()
#C训练
traj = t.cat((hist, fut), 0)
c_out = classified(traj)
loss_c = BCEloss(c_out, lat_enc)
optimizer_c.zero_grad()
loss_c.backward()
a = t.nn.utils.clip_grad_norm_(classified.parameters(), 10)
optimizer_c.step()
#D训练
real_data, _ = discriminator(traj, nbrs, mask)
g_out, _, _ = generator(hist, nbrs, mask, lat_enc, lon_enc, va,
nbrsva, lane, nbrslane, dis, nbrsdis)
fake_data, _ = discriminator(t.cat((hist, g_out), 0), nbrs, mask)
real_label = t.ones_like(real_data)
fake_label = t.zeros_like(fake_data)
loss_d1 = BCEloss(real_data, real_label)
loss_d2 = BCEloss(fake_data, fake_label)
loss_d = loss_d1 + loss_d2
optimizer_d.zero_grad()
loss_d.backward()
a = t.nn.utils.clip_grad_norm_(discriminator.parameters(), 10)
optimizer_d.step()
#G训练
g_out, lat_pred, lon_pred = generator(hist, nbrs, mask, lat_enc,
lon_enc, va, nbrsva, lane,
nbrslane, dis, nbrsdis)
loss_g1 = MSELoss2(g_out, fut, op_mask)
loss_gx = BCEloss(lat_pred, lat_enc) + BCEloss(lon_pred, lon_enc)
traj_fake = t.cat((hist, g_out), 0)
traj_true = t.cat((hist, fut), 0)
c_out = classified(traj_fake)
loss_g3 = BCEloss(c_out, lat_enc)
_, outp1 = discriminator(traj_fake, nbrs, mask)
_, outp2 = discriminator(traj_true, nbrs, mask)
loss_g4 = MSEloss(outp1, outp2)
loss_g = loss_g1 + loss_gx + 5 * loss_g3 + 5 * loss_g4
optimizer_g.zero_grad()
loss_g.backward()
a = t.nn.utils.clip_grad_norm_(generator.parameters(), 10)
optimizer_g.step()
loss_gi1 += loss_g1.item()
loss_gix += loss_gx.item()
loss_gi3 += loss_g3.item()
loss_gi4 += loss_g4.item()
if idx % 100 == 99:
print('mse:', loss_gi1 / 100, '|c1:', loss_gix / 100, '|c:',
loss_gi3 / 100, '|d:', loss_gi4 / 100)
file.write(str(loss_gi1 / 100) + ',')
loss_gi1 = 0
loss_gix = 0
loss_gi3 = 0
loss_gi4 = 0
avg_val_loss = 0
val_batch_count = 0
print('startval:')
for i, data in enumerate(valDataloader):
hist, nbrs, mask, lat_enc, lon_enc, fut, op_mask, va, nbrsva, lane, nbrslane, dis, nbrsdis = data
hist = hist.cuda()
nbrs = nbrs.cuda()
mask = mask.cuda()
lat_enc = lat_enc.cuda()
lon_enc = lon_enc.cuda()
fut = fut.cuda()
op_mask = op_mask.cuda()
va = va.cuda()
nbrsva = nbrsva.cuda()
lane = lane.cuda()
nbrslane = nbrslane.cuda()
dis = dis.cuda()
nbrsdis = nbrsdis.cuda()
fut_pred, _, _ = generator(hist, nbrs, mask, lat_enc, lon_enc, va,
nbrsva, lane, nbrslane, dis, nbrsdis)
l = MSELoss2(fut_pred, fut, op_mask)
avg_val_loss += l.item()
val_batch_count += 1
print('valmse:', avg_val_loss / val_batch_count)
t.save(generator.state_dict(),
'checkpoints/6f/epoch' + str(epoch + 1) + '_g.tar')
t.save(discriminator.state_dict(),
'checkpoints/6f/epoch' + str(epoch + 1) + '_d.tar')
t.save(classified.state_dict(),
'checkpoints/6f/epoch' + str(epoch + 1) + '_c.tar')
scheduler_g.step()
scheduler_d.step()
scheduler_c.step()
file.close()
def train_and_eval(self):
print("Training the TuckER model...")
self.entity_idxs = {d.entities[i]: i for i in range(len(d.entities))}
self.relation_idxs = {
d.relations[i]: i
for i in range(len(d.relations))
}
train_data_idxs = self.get_data_idxs(d.train_data)
print("Number of training data points: %d" % len(train_data_idxs))
model = TuckER(d, self.ent_vec_dim, self.rel_vec_dim, **self.kwargs)
if self.cuda:
model.cuda()
model.init()
opt = torch.optim.Adam(model.parameters(), lr=self.learning_rate)
if self.decay_rate:
scheduler = ExponentialLR(opt, self.decay_rate)
er_vocab = self.get_er_vocab(train_data_idxs)
er_vocab_pairs = list(er_vocab.keys())
print("Starting training...")
for it in range(1, self.num_iterations + 1):
start_train = time.time()
model.train()
losses = []
np.random.shuffle(er_vocab_pairs)
for j in range(0, len(er_vocab_pairs), self.batch_size):
data_batch, targets = self.get_batch(er_vocab, er_vocab_pairs,
j)
opt.zero_grad()
e1_idx = torch.tensor(data_batch[:, 0])
r_idx = torch.tensor(data_batch[:, 1])
if self.cuda:
e1_idx = e1_idx.cuda()
r_idx = r_idx.cuda()
predictions = model.forward(e1_idx, r_idx)
if self.label_smoothing:
targets = ((1.0 - self.label_smoothing) *
targets) + (1.0 / targets.size(1))
loss = model.loss(predictions, targets)
loss.backward()
opt.step()
losses.append(loss.item())
if self.decay_rate:
scheduler.step()
print("Iteration %d, time taken %.3f, loss %.10f " %
(it, time.time() - start_train, np.mean(losses)))
model.eval()
with torch.no_grad():
if not it % 50 and it <= 1500:
print("Validation:")
self.evaluate(model, d.valid_data)
print("Test:")
self.evaluate(model, d.test_data)
if not it % 1 and it > 1500:
print("Validation:")
self.evaluate(model, d.valid_data)
print("Test:")
self.evaluate(model, d.test_data)
def run(train_batch_size, val_batch_size, epochs, lr, log_interval, channels, classes):
# load model
model = get_featureExtractor('dense121')(classes, len(channels))
# load saved weights
if args.pretrained:
if len(args.checkpoint)>0:
cpfile = 'models/'+args.checkpoint+'.pth'
ep0 = epoch = int(args.checkpoint.rsplit('_',1)[1])
else:
cpfile = 'models/Model_pretrained_DenseNet121.pth'
ep0 = epoch = 1
checkpoint = torch.load(cpfile)
model.load_state_dict(checkpoint)
else:
ep0 = epoch = 0
# to gpu
if torch.cuda.is_available():
device = "cuda"
model.cuda(args.gpu)
else:
device = "cpu"
# augmentation
if args.augmentation:
aug = ImgAugTransform()
else:
aug = None
# load data generator
train_loader, val_loader = get_data_loaders(train_batch_size, val_batch_size, channels, classes, device=device) # aug=aug
# for parallel
if args.distributed:
model = nn.DataParallel(model)
# set optimizer and loss function
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss() #
# exponential decreasing learning rate
lr_scheduler = ExponentialLR(optimizer, gamma=0.95)
# process bar
desc = "ITERATION - loss: {:.2f}"
pbar = tqdm(
initial=0, leave=False, total=len(train_loader),
desc=desc.format(0)
)
# initialize parameters for iterating over epochs
burn=2 # epoch burn for
patience=3 # number of epochs to early stop training
vl_track=[] # tracking the validation loss
save_interval=1 # number of intervals to save weights
n_saved=5 # number of weights to keep
tlen = len(train_loader)
vlen = len(val_loader)
# print(tlen)
while epoch < epochs:
# frozen the pretrained layers, tran the fully connected classification layer
print(f'{epoch+1}/{epochs}')
print(f"Learning rate: {lr}")
if epoch == 0:
for name, child in model.named_children(): # module.
# pbar.log_message(name)
if name == 'fc':
print(name + ' is unfrozen')
for param in child.parameters():
param.requires_grad = True
else:
print(name + ' are frozen')
for param in child.parameters():
param.requires_grad = False
if epoch == burn:
print("Turn on all the layers")
# for name, child in model.named_children(): # module.
for param in model.features.parameters():
param.requires_grad = True
model.train()
# start of the epoch,
tloss = 0
acc = np.zeros(1)
t0 = time()
for i, (x, y) in enumerate(train_loader):
if aug:
x = torch.from_numpy(aug(x))
x = x.to(device)
y = torch.tensor(y).long().to(device)
t1 = time()
optimizer.zero_grad()
output = model(x)
# one hot for Binary Cross Entropy
# target = torch.zeros_like(output, device=device)
# target[np.arange(x.size(0)), y] = 1
loss = criterion(output, y)
loss.backward()
t2 = time()
optimizer.step()
t3 = time()
tloss += loss.item()
acc += accuracy(output, y)
del loss, output, y, x
torch.cuda.empty_cache()
if i>0 and i % log_interval == 0:
pbar.desc = desc.format(tloss/(i+1))
pbar.update(log_interval)
# print(t1-t0, t2-t1, t3-t2)
# print(psutil.cpu_percent())
# print(psutil.virtual_memory()) # physical memory usage
# done epoch
pbar.desc = desc.format(tloss/tlen)
pbar.update(tlen%log_interval)
# save checkpoints
if (epoch+1)%save_interval==0:
ch = ''.join([str(i) for i in channels])
torch.save(model.state_dict(), f'models/Model_pretrained_{ch}_DenseNet121_{epoch+1}.pth')
if (epoch+1)//save_interval>n_saved:
try:
os.remove(f'models/Model_pretrained_{ch}_DenseNet121_{epoch+1-save_interval*n_saved}.pth')
except:
pass
model.eval()
# compute loss and accuracy of validation
vloss = 0
vacc = np.zeros(1)
for i, (x, y) in enumerate(val_loader):
x = x.to(device)
y = torch.tensor(y).long().to(device)
optimizer.zero_grad()
output = model(x)
loss = criterion(output, y)
vloss += loss.item()
vacc += accuracy(output, y)
del loss, output, y, x
torch.cuda.empty_cache()
vl_track.append(vloss)
print('Epoch {} -> Train Loss: {:.4f}, ACC: {:.2f}%'.format(epoch+1, tloss/tlen, acc[0]/tlen))
print('Epoch {} -> Validation Loss: {:.4f}, ACC: {:.2f}%'.format(epoch+1, vloss/vlen, vacc[0]/vlen))
# reset process bar
pbar.desc = desc.format(tloss/(i+1))
pbar.update(log_interval)
pbar.n = pbar.last_print_n = 0
# stop training if vloss keeps increasing for patience
if epoch-ep0>=patience and all([vl_track[-1-i]>vl_track[-2-i] for i in range(patience-1)]):
break
# update learning
if epoch>=burn:
lr_scheduler.step()
lr = float(optimizer.param_groups[0]['lr'])
epoch += 1
# checkpoint ignite issue https://github.com/pytorch/ignite/pull/182
pbar.close()