def __init__(self, epochs, ID_txt, batch_size, n_cpu):
self.device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu")
self.epochs = epochs
self.model = DummyModel().to(self.device)
self.optimizer = torch.optim.Adam(self.model.parameters())
self.dataloader = TrainDataLoader(ID_txt,
batch_size=batch_size,
num_workers=n_cpu)
def main(args, model):
tr_dataset = TrainDataset(json_dir=args.json_path,
batch_size=args.batch_size)
cv_dataset = CvDataset(json_dir=args.json_path,
batch_size=args.cv_batch_size)
tr_loader = TrainDataLoader(data_set=tr_dataset,
batch_size=1,
num_workers=args.num_workers)
cv_loader = CvDataLoader(data_set=cv_dataset,
batch_size=1,
num_workers=args.num_workers)
data= {'tr_loader': tr_loader, 'cv_loader': cv_loader}
print(model)
# count the parameter number of the network
print('The number of trainable parameters of the net is:%d' % (numParams(model)))
model.cuda()
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=args.l2)
solver = Solver(data, model, optimizer, args)
solver.train()
line = line.replace("\n", "")
# rel_keys.append(line.split("\t")[0])
# print(ent_dic.get(line.split("#")[0]))
# print(rel_dic.get(line.split("#")[1]))
# print(ent_dic.get(line.split("#")[2]))
f.write(str(ent_dic.get(line.split("#")[0])) + "," + str(rel_dic.get(line.split("#")[1])) + "," + str(
ent_dic.get(line.split("#")[2])) + "\n")
f.close()
'''
# dataloader for training
train_dataloader = TrainDataLoader(in_path="./benchmarks/FB15K237/",
nbatches=100,
threads=8,
sampling_mode="normal",
bern_flag=1,
filter_flag=1,
neg_ent=25,
neg_rel=0)
# define the model
transh = TransH(ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=1,
p_norm=1,
norm_flag=True)
# change 1 transh_100
matrix = transh.load_checkpoint('./checkpoint/transh_100.ckpt')
print(transh.get_parameters("numpy").keys())
def main():
'''parameter initialization'''
args = parser.parse_args()
exp_name_dir = AverageMeter.experiment_name_dir(args.experiment_name)
'''model on gpu'''
model = TrackerSiamRPN()
'''setup train data loader'''
name = 'All'
assert name in ['VID', 'GOT-10k', 'All']
if name == 'GOT-10k':
root_dir = args.train_path
seq_dataset = GOT10k(root_dir, subset='train')
elif name == 'VID':
root_dir = '/home/arbi/desktop/ILSVRC2017_VID/ILSVRC'
seq_dataset = ImageNetVID(root_dir, subset=('train'))
elif name == 'All':
root_dir_vid = '/home/arbi/desktop/ILSVRC2017_VID/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('train'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='train')
seq_dataset = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset', len(seq_dataset))
train_data = TrainDataLoader(seq_dataset, name)
train_loader = DataLoader(dataset=train_data,
batch_size=1,
shuffle=True,
num_workers=16,
pin_memory=True)
'''setup val data loader'''
name = 'All'
assert name in ['VID', 'GOT-10k', 'All']
if name == 'GOT-10k':
root_dir = args.train_path
seq_dataset_val = GOT10k(root_dir, subset='val')
elif name == 'VID':
root_dir = '/home/arbi/desktop/ILSVRC2017_VID/ILSVRC'
seq_dataset_val = ImageNetVID(root_dir, subset=('val'))
elif name == 'All':
root_dir_vid = '/home/arbi/desktop/ILSVRC2017_VID/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('val'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='val')
seq_dataset_val = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset_val', len(seq_dataset_val))
val_data = TrainDataLoader(seq_dataset_val, name)
val_loader = DataLoader(dataset=val_data,
batch_size=1,
shuffle=True,
num_workers=16,
pin_memory=True)
'''load weights'''
init_weights(model)
if not args.checkpoint_path == None:
assert os.path.isfile(
args.checkpoint_path), '{} is not valid checkpoint_path'.format(
args.checkpoint_path)
try:
model.net.load_state_dict(
torch.load(args.checkpoint_path,
map_location=lambda storage, loc: storage))
print('You are loading the model.load_state_dict')
except:
init_weights(model)
'''train phase'''
closses, rlosses, tlosses, steps = AverageMeter(), AverageMeter(
), AverageMeter(), AverageMeter()
for epoch in range(config.epoches):
print('Train epoch {}/{}'.format(epoch + 1, config.epoches))
with tqdm(total=config.train_epoch_size) as progbar:
for i, dataset in enumerate(train_loader):
closs, rloss, loss, cur_lr = model.step(epoch,
dataset,
backward=True)
closs_ = closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
closses.update(closs.cpu().item())
rlosses.update(rloss.cpu().item())
tlosses.update(loss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(closses.avg),
rloss='{:05.3f}'.format(rlosses.avg),
tloss='{:05.3f}'.format(tlosses.avg))
progbar.update()
if i >= config.train_epoch_size - 1:
'''save plot'''
closses.closs_array.append(closses.avg)
rlosses.rloss_array.append(rlosses.avg)
tlosses.loss_array.append(tlosses.avg)
steps.update(steps.count)
steps.steps_array.append(steps.count)
steps.plot(exp_name_dir)
'''save model'''
model.save(model, exp_name_dir, epoch)
break
def main():
'''parameter initialization'''
args = parser.parse_args()
exp_name_dir = util.experiment_name_dir(args.experiment_name)
'''model on gpu'''
model = TrackerSiamRPN()
model.net.init_weights()
'''setup train data loader'''
name = 'VID'
assert name in ['VID', 'GOT-10k', 'All']
if name == 'GOT-10k':
root_dir = args.train_path
seq_dataset = GOT10k(root_dir, subset='val')
elif name == 'VID':
root_dir = '/home/arbi/desktop/ILSVRC2017_VID'
seq_dataset = ImageNetVID(root_dir, subset=('train', 'val'))
elif name == 'All':
root_dir_vid = '/home/arbi/desktop/ILSVRC2017_VID'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('train'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='train')
seq_dataset = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset', len(seq_dataset))
train_data = TrainDataLoader(seq_dataset, name)
train_loader = DataLoader(dataset=train_data,
batch_size=64,
shuffle=True,
num_workers=16,
pin_memory=True)
'''setup val data loader'''
name = 'GOT-10k'
assert name in ['VID', 'GOT-10k', 'All']
if name == 'GOT-10k':
root_dir = args.train_path
seq_dataset_val = GOT10k(root_dir, subset='val')
elif name == 'VID':
root_dir = '/home/arbi/desktop/ILSVRC2017_VID'
seq_dataset_val = ImageNetVID(root_dir, subset=('val'))
elif name == 'All':
root_dir_vid = '/home/arbi/desktop/ILSVRC2017_VID/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('val'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='val')
seq_dataset_val = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset_val', len(seq_dataset_val))
val_data = TrainDataLoader(seq_dataset_val, name)
val_loader = DataLoader(dataset=val_data,
batch_size=8,
shuffle=False,
num_workers=16,
pin_memory=True)
'''load weights'''
if not args.checkpoint_path == None:
assert os.path.isfile(
args.checkpoint_path), '{} is not valid checkpoint_path'.format(
args.checkpoint_path)
checkpoint = torch.load(args.checkpoint_path, map_location='cpu')
if 'model' in checkpoint.keys():
model.net.load_state_dict(
torch.load(args.checkpoint_path, map_location='cpu')['model'])
else:
model.net.load_state_dict(
torch.load(args.checkpoint_path, map_location='cpu'))
#model.net.load_state_dict(torch.load(args.checkpoint_path, map_location=lambda storage, loc: storage))
print('You are loading the model.load_state_dict')
elif config.pretrained_model:
#print("init with pretrained checkpoint %s" % config.pretrained_model + '\n')
#print('------------------------------------------------------------------------------------------------ \n')
checkpoint = torch.load(config.pretrained_model)
# change name and load parameters
checkpoint = {
k.replace('features.features', 'featureExtract'): v
for k, v in checkpoint.items()
}
model_dict = model.net.state_dict()
model_dict.update(checkpoint)
model.net.load_state_dict(model_dict)
torch.cuda.empty_cache()
'''train phase'''
train_closses, train_rlosses, train_tlosses = AverageMeter(), AverageMeter(
), AverageMeter()
val_closses, val_rlosses, val_tlosses = AverageMeter(), AverageMeter(
), AverageMeter()
train_val_plot = SavePlot(exp_name_dir, 'train_val_plot')
for epoch in range(config.epoches):
model.net.train()
if config.fix_former_3_layers:
if 1 > 1:
util.freeze_layers(model.net.module)
else:
util.freeze_layers(model.net)
print('Train epoch {}/{}'.format(epoch + 1, config.epoches))
train_loss = []
with tqdm(total=config.train_epoch_size) as progbar:
for i, dataset in enumerate(train_loader):
closs, rloss, loss = model.step(epoch, dataset, train=True)
closs_ = closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
train_closses.update(closs.cpu().item())
train_rlosses.update(rloss.cpu().item())
train_tlosses.update(loss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(train_closses.avg),
rloss='{:05.3f}'.format(train_rlosses.avg),
tloss='{:05.3f}'.format(train_tlosses.avg))
progbar.update()
train_loss.append(train_tlosses.avg)
if i >= config.train_epoch_size - 1:
'''save plot'''
#train_val_plot.update(train_tlosses.avg, train_label = 'total loss')
'''save model'''
model.save(model, exp_name_dir, epoch)
break
train_loss = np.mean(train_loss)
'''val phase'''
val_loss = []
with tqdm(total=config.val_epoch_size) as progbar:
print('Val epoch {}/{}'.format(epoch + 1, config.epoches))
for i, dataset in enumerate(val_loader):
val_closs, val_rloss, val_tloss = model.step(epoch,
dataset,
train=False)
closs_ = val_closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
val_closses.update(val_closs.cpu().item())
val_rlosses.update(val_rloss.cpu().item())
val_tlosses.update(val_tloss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(val_closses.avg),
rloss='{:05.3f}'.format(val_rlosses.avg),
tloss='{:05.3f}'.format(val_tlosses.avg))
progbar.update()
val_loss.append(val_tlosses.avg)
if i >= config.val_epoch_size - 1:
break
val_loss = np.mean(val_loss)
train_val_plot.update(train_loss, val_loss)
print('Train loss: {}, val loss: {}'.format(train_loss, val_loss))
def main():
'''parameter initialization'''
args = parser.parse_args()
exp_name_dir = util.experiment_name_dir(args.experiment_name)
'''model on gpu'''
model = TrackerSiamRPN()
'''setup train data loader'''
name = 'GOT-10k'
assert name in ['VID', 'GOT-10k', 'All', 'RGBT-234']
if name == 'GOT-10k':
root_dir_RGBT234 = args.train_path
root_dir_GTOT = '/home/krautsct/Grayscale-Thermal-Dataset'
seq_dataset_rgb = GOT10k(root_dir_RGBT234, subset='train_i')
seq_dataset_i = GOT10k(root_dir_RGBT234,
subset='train_i',
visible=False)
elif name == 'VID':
root_dir = '/home/arbi/desktop/ILSVRC'
seq_dataset = ImageNetVID(root_dir, subset=('train'))
elif name == 'All':
root_dir_vid = '/home/arbi/desktop/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('train'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='train')
seq_dataset = util.data_split(seq_datasetVID, seq_datasetGOT)
elif name == 'RGBT-234':
root_dir = args.train_path
seq_dataset = RGBTSequence(root_dir, subset='train')
seq_dataset_val = RGBTSequence(root_dir, subset='val')
print('seq_dataset', len(seq_dataset_rgb))
train_z_transforms = transforms.Compose([ToTensor()])
train_x_transforms = transforms.Compose([ToTensor()])
train_data_ir = TrainDataLoader_ir(seq_dataset_i, train_z_transforms,
train_x_transforms, name)
anchors = train_data_ir.anchors
train_loader_ir = DataLoader(dataset=train_data_ir,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=config.train_num_workers,
pin_memory=True)
train_data_rgb = TrainDataLoader(seq_dataset_rgb, train_z_transforms,
train_x_transforms, name)
anchors = train_data_rgb.anchors
train_loader_rgb = DataLoader(dataset=train_data_rgb,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=config.train_num_workers,
pin_memory=True)
'''setup val data loader'''
name = 'GOT-10k'
assert name in ['VID', 'GOT-10k', 'All', 'RGBT-234']
if name == 'GOT-10k':
val_dir = '/home/krautsct/RGB-t-Val'
seq_dataset_val_rgb = GOT10k(val_dir, subset='train_i')
seq_dataset_val_ir = GOT10k(val_dir, subset='train_i', visible=False)
elif name == 'VID':
root_dir = '/home/arbi/desktop/ILSVRC'
seq_dataset_val = ImageNetVID(root_dir, subset=('val'))
elif name == 'All':
root_dir_vid = '/home/arbi/desktop/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('val'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='val')
seq_dataset_val = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset_val', len(seq_dataset_val_rgb))
valid_z_transforms = transforms.Compose([ToTensor()])
valid_x_transforms = transforms.Compose([ToTensor()])
val_data = TrainDataLoader_ir(seq_dataset_val_ir, valid_z_transforms,
valid_x_transforms, name)
val_loader_ir = DataLoader(dataset=val_data,
batch_size=config.valid_batch_size,
shuffle=False,
num_workers=config.valid_num_workers,
pin_memory=True)
val_data_rgb = TrainDataLoader(seq_dataset_val_rgb, valid_z_transforms,
valid_x_transforms, name)
val_loader_rgb = DataLoader(dataset=val_data_rgb,
batch_size=config.valid_batch_size,
shuffle=False,
num_workers=config.valid_num_workers,
pin_memory=True)
val_losslist = []
'''load weights'''
if not args.checkpoint_path == None:
assert os.path.isfile(
args.checkpoint_path), '{} is not valid checkpoint_path'.format(
args.checkpoint_path)
checkpoint = torch.load(args.checkpoint_path, map_location='cpu')
if 'model' in checkpoint.keys():
model.net.load_state_dict(
torch.load(args.checkpoint_path, map_location='cpu')['model'])
else:
model.net.load_state_dict(
torch.load(args.checkpoint_path, map_location='cpu'))
torch.cuda.empty_cache()
print('You are loading the model.load_state_dict')
elif config.pretrained_model:
checkpoint = torch.load(config.pretrained_model)
# change name and load parameters
checkpoint = {
k.replace('features.features', 'featureExtract'): v
for k, v in checkpoint.items()
}
model_dict = model.net.state_dict()
model_dict.update(checkpoint)
model.net.load_state_dict(model_dict)
#torch.cuda.empty_cache()
'''train phase'''
train_closses, train_rlosses, train_tlosses = AverageMeter(), AverageMeter(
), AverageMeter()
val_closses, val_rlosses, val_tlosses = AverageMeter(), AverageMeter(
), AverageMeter()
#train_val_plot = SavePlot(exp_name_dir, 'train_val_plot')
val_plot = SavePlotVal(exp_name_dir, 'val_plot')
for epoch in range(config.epoches):
model.net.train()
if config.fix_former_3_layers:
util.freeze_layers(model.net)
print('Train epoch {}/{}'.format(epoch + 1, config.epoches))
train_loss = []
with tqdm(total=config.train_epoch_size) as progbar:
for i, (dataset_rgb, dataset_ir) in enumerate(
zip(train_loader_rgb, train_loader_ir)):
#for i, dataset_rgb in enumerate(train_loader_rgb):
closs, rloss, loss = model.step(epoch,
dataset_rgb,
dataset_ir,
anchors,
epoch,
i,
train=True)
closs_ = closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
train_closses.update(closs.cpu().item())
train_rlosses.update(rloss.cpu().item())
train_tlosses.update(loss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(train_closses.avg),
rloss='{:05.5f}'.format(train_rlosses.avg),
tloss='{:05.3f}'.format(train_tlosses.avg))
progbar.update()
train_loss.append(train_tlosses.avg)
if i >= config.train_epoch_size - 1:
'''save model'''
model.save(model, exp_name_dir, epoch)
break
train_loss = np.mean(train_loss)
'''val phase'''
val_loss = []
with tqdm(total=config.val_epoch_size) as progbar:
print('Val epoch {}/{}'.format(epoch + 1, config.epoches))
for i, (dataset_rgb,
dataset_ir) in enumerate(zip(val_loader_rgb,
val_loader_ir)):
#for i, dataset_rgb in enumerate(val_loader_rgb):
val_closs, val_rloss, val_tloss = model.step(epoch,
dataset_rgb,
dataset_ir,
anchors,
epoch,
train=False)
closs_ = val_closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
val_closses.update(val_closs.cpu().item())
val_rlosses.update(val_rloss.cpu().item())
val_tlosses.update(val_tloss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(val_closses.avg),
rloss='{:05.5f}'.format(val_rlosses.avg),
tloss='{:05.3f}'.format(val_tlosses.avg))
progbar.update()
val_loss.append(val_tlosses.avg)
if i >= config.val_epoch_size - 1:
break
val_loss = np.mean(val_loss)
#train_val_plot.update(train_loss, val_loss)
val_plot.update(val_loss)
val_losslist.append(val_loss)
print('Train loss: {}, val loss: {}'.format(train_loss, val_loss))
record_path = os.path.dirname(exp_name_dir)
if not os.path.isdir(record_path):
os.makedirs(record_path)
record_file = os.path.join(exp_name_dir, 'val_losses.txt')
np.savetxt(record_file, val_losslist, fmt='%.3f', delimiter=',')
def main():
args = parse_args()
print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) +
"Build training dataloader",
flush=True)
# dataloader for training
train_dataloader = TrainDataLoader(tri_file=args.link,
ent_file=args.node,
rel_file=args.rela,
batch_size=args.batch_size,
bern_flag=True,
filter_flag=True,
neg_ent=args.neg_ent,
neg_rel=args.neg_rel)
print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) +
"Define ComplEx model",
flush=True)
# define the model
complEx = ComplEx(
ent_tot=train_dataloader.get_ent_tot(),
rel_tot=train_dataloader.get_rel_tot(),
dim=args.dim // 2 # half real, half imaginary
)
print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) +
"Define loss function",
flush=True)
# define the loss function
model = NegativeSampling(model=complEx,
loss=SoftplusLoss(),
batch_size=train_dataloader.get_batch_size(),
regul_rate=args.regul_rate)
print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) +
"Start training",
flush=True)
# train the model
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=args.train_times,
alpha=args.alpha,
use_gpu=args.if_gpu,
opt_method=args.opt_method)
trainer.run()
print(time.strftime("%a, %d %b %Y %H:%M:%S +0000: ", time.localtime()) +
"Save the embeddings",
flush=True)
# save the embeddings
res = complEx.get_parameters("numpy")
with open(args.output, 'w') as file:
file.write(
f'size={args.dim}, regul_rate={args.regul_rate}, batch_size={args.batch_size}, train_times={args.train_times}, alpha={args.alpha}, opt_method={args.opt_method}\n'
)
for nid, (real_half, imag_half) in enumerate(
zip(res['ent_re_embeddings.weight'],
res['ent_im_embeddings.weight'])):
file.write('{}\t{}\n'.format(
nid, ' '.join(map(str, np.concatenate([real_half,
imag_half])))))
from config import Trainer, Tester
from module.model import Linear
from module.loss import RMSELoss
from module.strategy import NoStrategy
from data import TrainDataLoader, TestDataLoader
# dataloader for training
train_dataloader = TrainDataLoader(
in_path="../data/database/area_passenger_index.csv",
x_dim=31,
hour=0,
batch_size=126,
threads=8)
# dataloader for test
# test_dataloader = TestDataLoader("./benchmarks/WN18RR/", "link")
# define the model
linear = Linear(x_dim=train_dataloader.get_x_dim(),
hour=train_dataloader.get_hour())
# define the loss function
model = NoStrategy(model=linear,
loss=RMSELoss(),
batch_size=train_dataloader.get_batch_size(),
regul_rate=0.0)
# train the model
trainer = Trainer(model=model,
data_loader=train_dataloader,
train_times=1,
def main():
'''parameter initialization'''
args = parser.parse_args()
exp_name_dir = util.experiment_name_dir(args.experiment_name)
'''model on gpu'''
model = TrackerSiamRPN()
'''setup train data loader'''
name = 'GOT-10k'
assert name in ['VID', 'GOT-10k', 'All']
if name == 'GOT-10k':
root_dir = args.train_path
seq_dataset = GOT10k(root_dir, subset='train')
elif name == 'VID':
root_dir = '/store_ssd/ILSVRC'
seq_dataset = ImageNetVID(root_dir, subset=('train'))
elif name == 'All':
root_dir_vid = '/store_ssd/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('train'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='train')
seq_dataset = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset', len(seq_dataset))
train_z_transforms = transforms.Compose([ToTensor()])
train_x_transforms = transforms.Compose([
RandomCrop([config.detection_img_size, config.detection_img_size],
config.max_translate),
RandomScale(config.scale_resize),
ToTensor()
])
train_data = TrainDataLoader(seq_dataset, train_z_transforms,
train_x_transforms, name)
anchors = train_data.anchors
train_loader = DataLoader(dataset=train_data,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=config.train_num_workers,
pin_memory=True)
'''setup val data loader'''
name = 'GOT-10k'
assert name in ['VID', 'GOT-10k', 'All']
if name == 'GOT-10k':
root_dir = args.train_path
seq_dataset_val = GOT10k(root_dir, subset='val')
elif name == 'VID':
root_dir = '/store_ssd/ILSVRC'
seq_dataset_val = ImageNetVID(root_dir, subset=('val'))
elif name == 'All':
root_dir_vid = '/store_ssd/ILSVRC'
seq_datasetVID = ImageNetVID(root_dir_vid, subset=('val'))
root_dir_got = args.train_path
seq_datasetGOT = GOT10k(root_dir_got, subset='val')
seq_dataset_val = util.data_split(seq_datasetVID, seq_datasetGOT)
print('seq_dataset_val', len(seq_dataset_val))
valid_z_transforms = transforms.Compose([ToTensor()])
valid_x_transforms = transforms.Compose([ToTensor()])
val_data = TrainDataLoader(seq_dataset_val, valid_z_transforms,
valid_x_transforms, name)
val_loader = DataLoader(dataset=val_data,
batch_size=config.valid_batch_size,
shuffle=False,
num_workers=config.valid_num_workers,
pin_memory=True)
'''load weights'''
if not args.checkpoint_path == None and args.epoch_i > 0:
checkpoint_path = os.path.join(args.checkpoint_path,
'model_e{}.pth'.format(args.epoch_i))
assert os.path.isfile(
checkpoint_path), '{} is not valid checkpoint_path'.format(
checkpoint_path)
checkpoint = torch.load(checkpoint_path, map_location='cpu')
if 'model' in checkpoint.keys():
model.net.load_state_dict(
torch.load(checkpoint_path, map_location='cpu')['model'])
else:
model.net.load_state_dict(
torch.load(checkpoint_path, map_location='cpu'))
torch.cuda.empty_cache()
print('You are loading the model.load_state_dict')
elif config.pretrained_model:
checkpoint = torch.load(config.pretrained_model)
# change name and load parameters
checkpoint = {
k.replace('features.features', 'featureExtract'): v
for k, v in checkpoint.items()
}
model_dict = model.net.state_dict()
model_dict.update(checkpoint)
model.net.load_state_dict(model_dict)
#torch.cuda.empty_cache()
print('You are loading the pretrained model')
'''train phase'''
train_closses, train_rlosses, train_tlosses = AverageMeter(), AverageMeter(
), AverageMeter()
val_closses, val_rlosses, val_tlosses = AverageMeter(), AverageMeter(
), AverageMeter()
train_val_plot = SavePlot(exp_name_dir, 'train_val_plot')
model.adjust_lr(args.epoch_i)
for epoch in range(args.epoch_i, config.epoches):
model.net.train()
if config.fix_former_3_layers:
util.freeze_layers(model.net)
print('Train epoch {}/{}'.format(epoch + 1, config.epoches))
train_loss = []
with tqdm(total=config.train_epoch_size) as progbar:
for i, dataset in enumerate(train_loader):
closs, rloss, loss = model.step(epoch,
dataset,
anchors,
i,
train=True)
closs_ = closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
train_closses.update(closs.cpu().item())
train_rlosses.update(rloss.cpu().item())
train_tlosses.update(loss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(train_closses.avg),
rloss='{:05.5f}'.format(train_rlosses.avg),
tloss='{:05.3f}'.format(train_tlosses.avg))
progbar.update()
train_loss.append(train_tlosses.avg)
if i >= config.train_epoch_size - 1:
'''save model'''
model.save(model, exp_name_dir, epoch)
break
train_loss = np.mean(train_loss)
'''val phase'''
val_loss = []
with tqdm(total=config.val_epoch_size) as progbar:
print('Val epoch {}/{}'.format(epoch + 1, config.epoches))
for i, dataset in enumerate(val_loader):
val_closs, val_rloss, val_tloss = model.step(epoch,
dataset,
anchors,
train=False)
closs_ = val_closs.cpu().item()
if np.isnan(closs_):
sys.exit(0)
val_closses.update(val_closs.cpu().item())
val_rlosses.update(val_rloss.cpu().item())
val_tlosses.update(val_tloss.cpu().item())
progbar.set_postfix(closs='{:05.3f}'.format(val_closses.avg),
rloss='{:05.5f}'.format(val_rlosses.avg),
tloss='{:05.3f}'.format(val_tlosses.avg))
progbar.update()
val_loss.append(val_tlosses.avg)
if i >= config.val_epoch_size - 1:
break
val_loss = np.mean(val_loss)
train_val_plot.update(train_loss, val_loss)
print('Train loss: {}, val loss: {}'.format(train_loss, val_loss))
from baselines.transe.trainer import Trainer
from baselines.transe.tester import Tester
from baselines.transe.model import TransE
from baselines.transe.negative_sampling import NegativeSampling
from baselines.transe.loss import MarginLoss
from data import TrainDataLoader, TestDataLoader
# dataloader for training
train_dataloader = TrainDataLoader(in_path="./data/FB15K237/",
nbatches=100,
threads=8,
sampling_mode="normal",
bern_flag=1,
filter_flag=1,
neg_ent=25,
neg_rel=0)
# dataloader for test
test_dataloader = TestDataLoader("./data/FB15K237/", "link")
# define the model
transe = TransE(
ent_tot=train_dataloader.get_ent_tot(), # entity total
rel_tot=train_dataloader.get_rel_tot(), # relation total
dim=200,
p_norm=1,
norm_flag=True)
# define the loss function
model = NegativeSampling(model=transe,
loss=MarginLoss(margin=5.0),