class Trainer:
"""
trainer class
"""
def __init__(self, cfg: Namespace, data: Dataset):
"""
Args:
cfg: configuration
data: train dataset
"""
self.cfg = cfg
self.train, self.valid = data.split(0.8)
RATING_FIELD.build_vocab(self.train)
self.device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu') # pylint: disable=no-member
self.batch_size = cfg.batch_size
if torch.cuda.is_available():
self.batch_size *= torch.cuda.device_count()
self.trn_itr = BucketIterator(
self.train,
device=self.device,
batch_size=self.batch_size,
shuffle=True,
train=True,
sort_within_batch=True,
sort_key=lambda exam: -len(exam.comment_text))
self.vld_itr = BucketIterator(
self.valid,
device=self.device,
batch_size=self.batch_size,
shuffle=False,
train=False,
sort_within_batch=True,
sort_key=lambda exam: -len(exam.comment_text))
self.log_step = 1000
if len(self.vld_itr) < 100:
self.log_step = 10
elif len(self.vld_itr) < 1000:
self.log_step = 100
bert_path = cfg.bert_path if cfg.bert_path else 'bert-base-cased'
self.model = BertForSequenceClassification.from_pretrained(
bert_path, num_labels=2)
pos_weight = (
len([exam for exam in self.train.examples if exam.target < 0.5]) /
len([exam for exam in self.train.examples if exam.target >= 0.5]))
pos_wgt_tensor = torch.tensor([1.0, pos_weight], device=self.device) # pylint: disable=not-callable
self.criterion = nn.CrossEntropyLoss(weight=pos_wgt_tensor)
if torch.cuda.is_available():
self.model = DataParallelModel(self.model.cuda())
self.criterion = DataParallelCriterion(self.criterion)
self.optimizer = optim.Adam(self.model.parameters(), cfg.learning_rate)
def run(self):
"""
do train
"""
max_f_score = -9e10
max_epoch = -1
for epoch in range(self.cfg.epoch):
train_loss = self._train_epoch(epoch)
metrics = self._evaluate(epoch)
max_f_score_str = f' < {max_f_score:.2f}'
if metrics['f_score'] > max_f_score:
max_f_score_str = ' is max'
max_f_score = metrics['f_score']
max_epoch = epoch
torch.save(self.model.state_dict(), self.cfg.model_path)
logging.info('EPOCH[%d]: train loss: %.6f, valid loss: %.6f, acc: %.2f,' \
' F: %.2f%s', epoch, train_loss, metrics['loss'],
metrics['accuracy'], metrics['f_score'], max_f_score_str)
if (epoch - max_epoch) >= self.cfg.patience:
logging.info('early stopping...')
break
logging.info('epoch: %d, f-score: %.2f', max_epoch, max_f_score)
def _train_epoch(self, epoch: int) -> float:
"""
train single epoch
Args:
epoch: epoch number
Returns:
average loss
"""
self.model.train()
progress = tqdm(self.trn_itr,
f'EPOCH[{epoch}]',
mininterval=1,
ncols=100)
losses = []
for step, batch in enumerate(progress, start=1):
outputs = self.model(batch.comment_text)
# output of model wrapped with DataParallelModel is a list of outputs from each GPU
# make input of DataParallelCriterion as a list of tuples
if isinstance(self.model, DataParallelModel):
loss = self.criterion([(output, ) for output in outputs],
batch.target)
else:
loss = self.criterion(outputs, batch.target)
losses.append(loss.item())
if step % self.log_step == 0:
avg_loss = sum(losses) / len(losses)
progress.set_description(f'EPOCH[{epoch}] ({avg_loss:.6f})')
loss.backward()
self.optimizer.step()
self.optimizer.zero_grad()
return sum(losses) / len(losses)
def _evaluate(self, epoch: int) -> Dict[str, float]:
"""
evaluate on validation data
Args:
epoch: epoch number
Returns:
metrics
"""
self.model.eval()
progress = tqdm(self.vld_itr,
f' EVAL[{epoch}]',
mininterval=1,
ncols=100)
losses = []
preds = []
golds = []
for step, batch in enumerate(progress, start=1):
with torch.no_grad():
outputs = self.model(batch.comment_text)
if isinstance(self.model, DataParallelModel):
loss = self.criterion([(output, ) for output in outputs],
batch.target)
for output in outputs:
preds.extend([(0 if o[0] < o[1] else 1)
for o in output])
else:
loss = self.criterion(outputs, batch.target)
preds.extend([(0 if output[0] < output[1] else 1)
for output in outputs])
losses.append(loss.item())
golds.extend([gold.item() for gold in batch.target])
if step % self.log_step == 0:
avg_loss = sum(losses) / len(losses)
progress.set_description(
f' EVAL[{epoch}] ({avg_loss:.6f})')
metrics = self._get_metrics(preds, golds)
metrics['loss'] = sum(losses) / len(losses)
return metrics
@classmethod
def _get_metrics(cls, preds: List[float],
golds: List[float]) -> Dict[str, float]:
"""
get metric values
Args:
preds: predictions
golds: gold standards
Returns:
metric
"""
assert len(preds) == len(golds)
true_pos = 0
false_pos = 0
false_neg = 0
true_neg = 0
for pred, gold in zip(preds, golds):
if pred >= 0.5:
if gold >= 0.5:
true_pos += 1
else:
false_pos += 1
else:
if gold >= 0.5:
false_neg += 1
else:
true_neg += 1
accuracy = (true_pos + true_neg) / (true_pos + false_pos + false_neg +
true_neg)
precision = 0.0
if (true_pos + false_pos) > 0:
precision = true_pos / (true_pos + false_pos)
recall = 0.0
if (true_pos + false_neg) > 0:
recall = true_pos / (true_pos + false_neg)
f_score = 0.0
if (precision + recall) > 0.0:
f_score = 2.0 * precision * recall / (precision + recall)
return {
'accuracy': 100.0 * accuracy,
'precision': 100.0 * precision,
'recall': 100.0 * recall,
'f_score': 100.0 * f_score,
}
def main_tr(args, crossVal):
dataLoad = ld.LoadData(args.data_dir, args.classes)
data = dataLoad.processData(crossVal, args.data_name)
# load the model
model = net.MiniSeg(args.classes, aux=True)
if not osp.isdir(osp.join(args.savedir + '_mod' + str(args.max_epochs))):
os.mkdir(args.savedir + '_mod' + str(args.max_epochs))
if not osp.isdir(
osp.join(args.savedir + '_mod' + str(args.max_epochs),
args.data_name)):
os.mkdir(
osp.join(args.savedir + '_mod' + str(args.max_epochs),
args.data_name))
saveDir = args.savedir + '_mod' + str(
args.max_epochs) + '/' + args.data_name + '/' + args.model_name
# create the directory if not exist
if not osp.exists(saveDir):
os.mkdir(saveDir)
if args.gpu and torch.cuda.device_count() > 1:
#model = torch.nn.DataParallel(model)
model = DataParallelModel(model)
if args.gpu:
model = model.cuda()
total_paramters = sum([np.prod(p.size()) for p in model.parameters()])
print('Total network parameters: ' + str(total_paramters))
# define optimization criteria
weight = torch.from_numpy(
data['classWeights']) # convert the numpy array to torch
if args.gpu:
weight = weight.cuda()
criteria = CrossEntropyLoss2d(weight, args.ignore_label) #weight
if args.gpu and torch.cuda.device_count() > 1:
criteria = DataParallelCriterion(criteria)
if args.gpu:
criteria = criteria.cuda()
# compose the data with transforms
trainDataset_main = myTransforms.Compose([
myTransforms.Normalize(mean=data['mean'], std=data['std']),
myTransforms.Scale(args.width, args.height),
myTransforms.RandomCropResize(int(32. / 1024. * args.width)),
myTransforms.RandomFlip(),
myTransforms.ToTensor()
])
trainDataset_scale1 = myTransforms.Compose([
myTransforms.Normalize(mean=data['mean'], std=data['std']),
myTransforms.Scale(int(args.width * 1.5), int(args.height * 1.5)),
myTransforms.RandomCropResize(int(100. / 1024. * args.width)),
myTransforms.RandomFlip(),
myTransforms.ToTensor()
])
trainDataset_scale2 = myTransforms.Compose([
myTransforms.Normalize(mean=data['mean'], std=data['std']),
myTransforms.Scale(int(args.width * 1.25), int(args.height * 1.25)),
myTransforms.RandomCropResize(int(100. / 1024. * args.width)),
myTransforms.RandomFlip(),
myTransforms.ToTensor()
])
trainDataset_scale3 = myTransforms.Compose([
myTransforms.Normalize(mean=data['mean'], std=data['std']),
myTransforms.Scale(int(args.width * 0.75), int(args.height * 0.75)),
myTransforms.RandomCropResize(int(32. / 1024. * args.width)),
myTransforms.RandomFlip(),
myTransforms.ToTensor()
])
valDataset = myTransforms.Compose([
myTransforms.Normalize(mean=data['mean'], std=data['std']),
myTransforms.Scale(args.width, args.height),
myTransforms.ToTensor()
])
# since we training from scratch, we create data loaders at different scales
# so that we can generate more augmented data and prevent the network from overfitting
trainLoader = torch.utils.data.DataLoader(myDataLoader.Dataset(
data['trainIm'], data['trainAnnot'], transform=trainDataset_main),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
trainLoader_scale1 = torch.utils.data.DataLoader(
myDataLoader.Dataset(data['trainIm'],
data['trainAnnot'],
transform=trainDataset_scale1),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
trainLoader_scale2 = torch.utils.data.DataLoader(
myDataLoader.Dataset(data['trainIm'],
data['trainAnnot'],
transform=trainDataset_scale2),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
trainLoader_scale3 = torch.utils.data.DataLoader(
myDataLoader.Dataset(data['trainIm'],
data['trainAnnot'],
transform=trainDataset_scale3),
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
pin_memory=True,
drop_last=True)
valLoader = torch.utils.data.DataLoader(myDataLoader.Dataset(
data['valIm'], data['valAnnot'], transform=valDataset),
batch_size=args.batch_size,
shuffle=False,
num_workers=args.num_workers,
pin_memory=True)
max_batches = len(trainLoader) + len(trainLoader_scale1) + len(
trainLoader_scale2) + len(trainLoader_scale3)
if args.gpu:
cudnn.benchmark = True
start_epoch = 0
if args.pretrained is not None:
state_dict = torch.load(args.pretrained)
new_keys = []
new_values = []
for idx, key in enumerate(state_dict.keys()):
if 'pred' not in key:
new_keys.append(key)
new_values.append(list(state_dict.values())[idx])
new_dict = OrderedDict(list(zip(new_keys, new_values)))
model.load_state_dict(new_dict, strict=False)
print('pretrained model loaded')
if args.resume is not None:
if osp.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
start_epoch = checkpoint['epoch']
args.lr = checkpoint['lr']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
log_file = osp.join(saveDir, 'trainValLog_' + args.model_name + '.txt')
if osp.isfile(log_file):
logger = open(log_file, 'a')
else:
logger = open(log_file, 'w')
logger.write("Parameters: %s" % (str(total_paramters)))
logger.write("\n%s\t%s\t\t%s\t%s\t%s\t%s\tlr" %
('CrossVal', 'Epoch', 'Loss(Tr)', 'Loss(val)',
'mIOU (tr)', 'mIOU (val)'))
logger.flush()
optimizer = torch.optim.Adam(model.parameters(),
args.lr, (0.9, 0.999),
eps=1e-08,
weight_decay=1e-4)
maxmIOU = 0
maxEpoch = 0
print(args.model_name + '-CrossVal: ' + str(crossVal + 1))
for epoch in range(start_epoch, args.max_epochs):
# train for one epoch
cur_iter = 0
train(args, trainLoader_scale1, model, criteria, optimizer, epoch,
max_batches, cur_iter)
cur_iter += len(trainLoader_scale1)
train(args, trainLoader_scale2, model, criteria, optimizer, epoch,
max_batches, cur_iter)
cur_iter += len(trainLoader_scale2)
train(args, trainLoader_scale3, model, criteria, optimizer, epoch,
max_batches, cur_iter)
cur_iter += len(trainLoader_scale3)
lossTr, overall_acc_tr, per_class_acc_tr, per_class_iu_tr, mIOU_tr, lr = \
train(args, trainLoader, model, criteria, optimizer, epoch, max_batches, cur_iter)
# evaluate on validation set
lossVal, overall_acc_val, per_class_acc_val, per_class_iu_val, mIOU_val = \
val(args, valLoader, model, criteria)
torch.save(
{
'epoch': epoch + 1,
'arch': str(model),
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
'lossTr': lossTr,
'lossVal': lossVal,
'iouTr': mIOU_tr,
'iouVal': mIOU_val,
'lr': lr
},
osp.join(
saveDir, 'checkpoint_' + args.model_name + '_crossVal' +
str(crossVal + 1) + '.pth.tar'))
# save the model also
model_file_name = osp.join(
saveDir, 'model_' + args.model_name + '_crossVal' +
str(crossVal + 1) + '_' + str(epoch + 1) + '.pth')
torch.save(model.state_dict(), model_file_name)
logger.write(
"\n%d\t\t%d\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.4f\t\t%.7f" %
(crossVal + 1, epoch + 1, lossTr, lossVal, mIOU_tr, mIOU_val, lr))
logger.flush()
print("\nEpoch No. %d:\tTrain Loss = %.4f\tVal Loss = %.4f\t mIOU(tr) = %.4f\t mIOU(val) = %.4f\n" \
% (epoch + 1, lossTr, lossVal, mIOU_tr, mIOU_val))
if mIOU_val >= maxmIOU:
maxmIOU = mIOU_val
maxEpoch = epoch + 1
torch.cuda.empty_cache()
logger.flush()
logger.close()
return maxEpoch, maxmIOU
