def run(dataset_train, dataset_dev, dataset_test, model_type, word_embed_size,
hidden_size, batch_size, use_cuda, n_epochs):
if model_type == 'base':
model = Baseline(vocab=dataset_train.vocab,
word_embed_size=word_embed_size,
hidden_size=hidden_size,
use_cuda=use_cuda,
inference=False)
else:
raise NotImplementedError
if use_cuda:
model = model.cuda()
optim_params = model.parameters()
optimizer = optim.Adam(optim_params, lr=10**-3)
print('start training')
for epoch in range(n_epochs):
train_loss, tokens, preds, golds = train(dataset_train, model,
optimizer, batch_size, epoch,
Phase.TRAIN, use_cuda)
dev_loss, tokens, preds, golds = train(dataset_dev, model, optimizer,
batch_size, epoch, Phase.DEV,
use_cuda)
logger = '\t'.join([
'epoch {}'.format(epoch + 1),
'TRAIN Loss: {:.9f}'.format(train_loss),
'DEV Loss: {:.9f}'.format(dev_loss)
])
print('\r' + logger, end='')
test_loss, tokens, preds, golds = train(dataset_test, model, optimizer,
batch_size, epoch, Phase.TEST,
use_cuda)
print('====', 'TEST', '=====')
print_scores(preds, golds)
output_results(tokens, preds, golds)
class Session:
def __init__(self):
self.log_dir = settings.log_dir
self.model_dir = settings.model_dir
ensure_dir(settings.log_dir)
ensure_dir(settings.model_dir)
logger.info('set log dir as %s' % settings.log_dir)
logger.info('set model dir as %s' % settings.model_dir)
##################################### Import models ###########################
self.feature_generator = Baseline(
last_stride=1, model_path=settings.pretrained_model_path)
self.feature_embedder_rgb = FeatureEmbedder(2048)
self.feature_embedder_ir = FeatureEmbedder(2048)
self.id_classifier = IdClassifier()
if torch.cuda.is_available():
self.feature_generator.cuda()
self.feature_embedder_rgb.cuda()
self.feature_embedder_ir.cuda()
self.id_classifier.cuda()
self.feature_generator = nn.DataParallel(self.feature_generator,
device_ids=range(
settings.num_gpu))
self.feature_embedder_rgb = nn.DataParallel(self.feature_embedder_rgb,
device_ids=range(
settings.num_gpu))
self.feature_embedder_ir = nn.DataParallel(self.feature_embedder_ir,
device_ids=range(
settings.num_gpu))
self.id_classifier = nn.DataParallel(self.id_classifier,
device_ids=range(
settings.num_gpu))
############################# Get Losses & Optimizers #########################
self.criterion_at = expATLoss()
self.criterion_identity = CrossEntropyLabelSmoothLoss(
settings.num_classes, epsilon=0.1) #torch.nn.CrossEntropyLoss()
opt_models = [
self.feature_generator, self.feature_embedder_rgb,
self.feature_embedder_ir, self.id_classifier
]
def make_optimizer(opt_models):
train_params = []
for opt_model in opt_models:
for key, value in opt_model.named_parameters():
if not value.requires_grad:
continue
lr = settings.BASE_LR
weight_decay = settings.WEIGHT_DECAY
if "bias" in key:
lr = settings.BASE_LR * settings.BIAS_LR_FACTOR
weight_decay = settings.WEIGHT_DECAY_BIAS
train_params += [{
"params": [value],
"lr": lr,
"weight_decay": weight_decay
}]
optimizer = torch.optim.Adam(train_params)
return optimizer
self.optimizer_G = make_optimizer(opt_models)
self.epoch_count = 0
self.step = 0
self.save_steps = settings.save_steps
self.num_workers = settings.num_workers
self.writers = {}
self.dataloaders = {}
self.sche_G = solver.WarmupMultiStepLR(self.optimizer_G,
milestones=settings.iter_sche,
gamma=0.1) # default setting
def tensorboard(self, name):
self.writers[name] = SummaryWriter(
os.path.join(self.log_dir, name + '.events'))
return self.writers[name]
def write(self, name, out):
for k, v in out.items():
self.writers[name].add_scalar(name + '/' + k, v, self.step)
out['G_lr'] = self.optimizer_G.param_groups[0]['lr']
out['step'] = self.step
out['eooch_count'] = self.epoch_count
outputs = ["{}:{:.4g}".format(k, v) for k, v in out.items()]
logger.info(name + '--' + ' '.join(outputs))
def save_checkpoints(self, name):
ckp_path = os.path.join(self.model_dir, name)
obj = {
'feature_generator': self.feature_generator.state_dict(),
'feature_embedder_rgb': self.feature_embedder_rgb.state_dict(),
'feature_embedder_ir': self.feature_embedder_ir.state_dict(),
'id_classifier': self.id_classifier.state_dict(),
'clock': self.step,
'epoch_count': self.epoch_count,
'opt_G': self.optimizer_G.state_dict(),
}
torch.save(obj, ckp_path)
def load_checkpoints(self, name):
ckp_path = os.path.join(self.model_dir, name)
try:
obj = torch.load(ckp_path)
print('load checkpoint: %s' % ckp_path)
except FileNotFoundError:
return
self.feature_generator.load_state_dict(obj['feature_generator'])
self.feature_embedder_rgb.load_state_dict(obj['feature_embedder_rgb'])
self.feature_embedder_ir.load_state_dict(obj['feature_embedder_ir'])
self.id_classifier.load_state_dict(obj['id_classifier'])
self.optimizer_G.load_state_dict(obj['opt_G'])
self.step = obj['clock']
self.epoch_count = obj['epoch_count']
self.sche_G.last_epoch = self.step
def load_checkpoints_delf_init(self, name):
ckp_path = os.path.join(self.model_dir, name)
obj = torch.load(ckp_path)
self.backbone.load_state_dict(obj['backbone'])
def cal_fea(self, x, domain_mode):
feat = self.feature_generator(x)
if domain_mode == 'rgb':
return self.feature_embedder_rgb(feat)
elif domain_mode == 'ir':
return self.feature_embedder_ir(feat)
def inf_batch(self, batch):
alpha = settings.alpha
beta = settings.beta
anchor_rgb, positive_rgb, negative_rgb, anchor_ir, positive_ir, \
negative_ir, anchor_label, modality_rgb, modality_ir = batch
if torch.cuda.is_available():
anchor_rgb = anchor_rgb.cuda()
positive_rgb = positive_rgb.cuda()
negative_rgb = negative_rgb.cuda()
anchor_ir = anchor_ir.cuda()
positive_ir = positive_ir.cuda()
negative_ir = negative_ir.cuda()
anchor_label = anchor_label.cuda()
anchor_rgb_features = self.cal_fea(anchor_rgb, 'rgb')
positive_rgb_features = self.cal_fea(positive_rgb, 'rgb')
negative_rgb_features = self.cal_fea(negative_rgb, 'rgb')
anchor_ir_features = self.cal_fea(anchor_ir, 'ir')
positive_ir_features = self.cal_fea(positive_ir, 'ir')
negative_ir_features = self.cal_fea(negative_ir, 'ir')
at_loss_rgb = self.criterion_at.forward(anchor_rgb_features,
positive_ir_features,
negative_ir_features)
at_loss_ir = self.criterion_at.forward(anchor_ir_features,
positive_rgb_features,
negative_rgb_features)
at_loss = at_loss_rgb + at_loss_ir
predicted_id_rgb = self.id_classifier(anchor_rgb_features)
predicted_id_ir = self.id_classifier(anchor_ir_features)
identity_loss = self.criterion_identity(predicted_id_rgb, anchor_label) + \
self.criterion_identity(predicted_id_ir, anchor_label)
loss_G = alpha * at_loss + beta * identity_loss
self.optimizer_G.zero_grad()
loss_G.backward()
self.optimizer_G.step()
self.write('train_stats', {
'loss_G': loss_G,
'at_loss': at_loss,
'identity_loss': identity_loss
})
valid_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=BATCH_SIZE,
num_workers=4,
collate_fn=collate_fn,
sampler=valid_sampler)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
batch_size=1,
num_workers=4,
shuffle=False)
config = {
"epochs": 100,
"device": get_device(),
"sampling": True,
"temperature": 1.0,
"max_sentence_length": 18
}
embedding_dim = 256
hidden_dim = 512
vocab_size = len(vocab)
model = Baseline(embedding_dim, hidden_dim, vocab_size, vanilla=False)
criterion = nn.CrossEntropyLoss()
optimizer = Adam(model.parameters(), lr=5e-4)
model.cuda()
train(model, optimizer, criterion, train_loader, valid_loader, vocab, config)
test(model, criterion, test_loader, vocab, config)
def main():
saver = utils.Saver(opt)
# randomize seed
opt.manualSeed = random.randint(1, 10000) # fix seed
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
torch.cuda.manual_seed_all(opt.manualSeed)
# load data
root = "data/modelnet40_ply_hdf5_2048/" #"data/modelnet40_normal_resampled"#
use_cuda = torch.cuda.is_available()
transforms_list = []
random_permute = utils.Random_permute(opt.num_points, delta=opt.distance)
# load transformations
if opt.random_input:
print("random_input")
transforms_list.append(random_permute)
# Load dataset / data loader
train_dataset = data.ModelNetDataset(
root,
train=True,
sort=opt.sort,
transform=transforms.Compose(transforms_list),
distance=opt.distance,
normal=opt.normal)
train_loader = DataLoader(train_dataset,
batch_size=opt.batchSize,
shuffle=True,
num_workers=opt.workers)
test_dataset = data.ModelNetDataset(root,
train=False,
sort=opt.sort,
distance=opt.distance,
normal=opt.normal)
test_loader = DataLoader(test_dataset,
batch_size=opt.batchSize,
shuffle=False,
num_workers=opt.workers)
# define model
ndim = 6 if opt.distance or opt.normal else 3
if opt.model == 'lstm':
model = Baseline(input_dim=ndim, maxout=opt.elem_max)
elif opt.model == 'lstm_mlp':
model = LSTM_mlp(input_dim=ndim,
maxout=opt.elem_max,
mlp=[64, 128, 256, 512],
fc=[512, 256, 40])
elif opt.model == 'test':
model = Test(input_dim=ndim, maxout=opt.elem_max)
# load speicified pre-trained model
if opt.path != '':
model.load_state_dict(torch.load(opt.path))
# define optimizer and loss function
optimizer = optim.Adam(model.parameters(),
lr=opt.learning_rate,
weight_decay=1e-5)
criterion = nn.CrossEntropyLoss()
# transfer model and criterion to cuda if exist
if use_cuda:
model = model.cuda(
) #nn.DataParallel(model).cuda()#model.cuda() #nn.DataParallel(model).cuda()
criterion = criterion.cuda()
best_model_wts = model.state_dict()
early_stopping = utils.Early_stopping(opt.early_stopping, patience=15)
saver.log_parameters(model.parameters())
for epoch in range(opt.nepoch):
adjust_learning_rate(optimizer, epoch, saver)
train(model, optimizer, criterion, saver, train_loader, epoch)
test_loss = test(model, criterion, saver, test_loader, epoch)
early_stopping.update(test_loss)
if early_stopping.stop():
break
saver.save_result()
