def train(config, device, train_loader, epoch):
"""
train an epoch
Args:
config (EasyDict): configurations for training
device (torch.device): the GPU or CPU used for training
train_loader (torch.utils.data.DataLoader): a DataLoader instance object for training set
epoch (int): current epoch
Returns:
err (float): error rate
"""
losses = AverageMeter()
evaluator = Evaluator(config.num_classes)
model.train()
with tqdm(train_loader) as pbar:
pbar.set_description('Train Epoch {}'.format(epoch))
for step, (input_, target) in enumerate(train_loader):
# move data to device
input_ = torch.tensor(input_, device=device, dtype=torch.float32)
target = torch.tensor(target, device=device, dtype=torch.long)
# forward and compute loss
output = model(input_)
loss = criterion(output, target)
# backward and update params
optimizer.zero_grad()
loss.backward()
optimizer.step()
# record loss and show it in the pbar
losses.update(loss.item(), input_.size(0))
postfix = OrderedDict({
'batch_loss': f'{losses.val:6.4f}',
'running_loss': f'{losses.avg:6.4f}'
})
pbar.set_postfix(ordered_dict=postfix)
pbar.update()
# visualization with TensorBoard
total_iter = (epoch - 1) * len(train_loader) + step + 1
writer.add_scalar('training_loss', losses.val, total_iter)
# update confusion matrix
true = target.cpu().numpy()
pred = output.max(dim=1)[1].cpu().numpy()
evaluator.update_matrix(true, pred)
return evaluator.error()
def train(self, epoch, train_loader):
"""
使用训练集训练一个epoch
Args:
epoch: (int) 第几代训练
train_loader: (torch.utils.data.DataLoader) 训练数据加载器
Returns:
err: (float) error rate
"""
losses = AverageMeter()
evaluator = Evaluator(self.cfgs.num_classes)
self.model.train()
with tqdm(train_loader) as pbar:
pbar.set_description('Train Epoch {}'.format(epoch))
for step, (input_, target) in enumerate(train_loader):
# move data to device
input_ = torch.tensor(input_, device=self.device, dtype=torch.float32)
target = torch.tensor(target, device=self.device, dtype=torch.long)
# forward and compute loss
output = self.model(input_)
loss = self.criterion(output, target)
# backward and update params
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
# record loss and show it in the pbar
losses.update(loss.item(), input_.size(0))
postfix = OrderedDict({'batch_loss': f'{losses.val:6.4f}', 'running_loss': f'{losses.avg:6.4f}'})
pbar.set_postfix(ordered_dict=postfix)
pbar.update()
# visualization with TensorBoard
total_iter = (epoch - 1) * len(train_loader) + step + 1
self.writer.add_scalar('training_loss', losses.val, total_iter)
# update confusion matrix
true = target.cpu().numpy()
pred = output.max(dim=1)[1].cpu().numpy()
evaluator.update_matrix(true, pred)
return evaluator.error()
def validate(config, device, val_loader, epoch):
"""
validate the model
Args:
config (EasyDict): configurations for training
device (torch.device): the GPU or CPU used for training
val_loader (torch.utils.data.DataLoader): a DataLoader instance object for validation set
epoch (int): current epoch
Returns:
err: (float) error rate
"""
losses = AverageMeter()
evaluator = Evaluator(config.num_classes)
model.eval()
with tqdm(val_loader) as pbar:
pbar.set_description('Valid Epoch {}'.format(epoch))
for i, (input_, target) in enumerate(val_loader):
# move data to GPU
input_ = torch.tensor(input_, device=device, dtype=torch.float32)
target = torch.tensor(target, device=device, dtype=torch.long)
with torch.no_grad():
# compute output and loss
output = model(input_)
loss = criterion(output, target)
# record loss and show it in the pbar
losses.update(loss.item(), input_.size(0))
postfix = OrderedDict({
'batch_loss': f'{losses.val:6.4f}',
'running_loss': f'{losses.avg:6.4f}'
})
pbar.set_postfix(ordered_dict=postfix)
pbar.update()
# update confusion matrix
true = target.cpu().numpy()
pred = output.max(dim=1)[1].cpu().numpy()
evaluator.update_matrix(true, pred)
return evaluator.error()
def validate(self, epoch, valid_loader):
"""
使用验证集测试模型的效果
Args:
epoch: (int) 第几代训练
valid_loader: (torch.utils.data.DataLoader) 验证数据加载器
Returns:
err: (float) error rate
"""
losses = AverageMeter()
evaluator = Evaluator(self.cfgs.num_classes)
self.model.eval()
with tqdm(valid_loader) as pbar:
pbar.set_description('Valid Epoch {}'.format(epoch))
for i, (input_, target) in enumerate(valid_loader):
# move data to GPU
input_ = torch.tensor(input_, device=self.device, dtype=torch.float32)
target = torch.tensor(target, device=self.device, dtype=torch.long)
with torch.no_grad():
# compute output and loss
output = self.model(input_)
loss = self.criterion(output, target)
# record loss and show it in the pbar
losses.update(loss.item(), input_.size(0))
postfix = OrderedDict({'batch_loss': f'{losses.val:6.4f}', 'running_loss': f'{losses.avg:6.4f}'})
pbar.set_postfix(ordered_dict=postfix)
pbar.update()
# update confusion matrix
true = target.cpu().numpy()
pred = output.max(dim=1)[1].cpu().numpy()
evaluator.update_matrix(true, pred)
return evaluator.error()
def test(device, model, test_loader, vis_conf_mat=False):
evaluator = Evaluator(Config.num_classes)
model.eval()
with tqdm(test_loader) as pbar:
pbar.set_description('Eval in test set')
for i, (input_, target) in enumerate(test_loader):
input_ = torch.tensor(input_, device=device, dtype=torch.float32)
target = torch.tensor(target, device=device, dtype=torch.long)
with torch.no_grad():
output = model(input_)
true = target.cpu().numpy()
pred = output.max(dim=1)[1].cpu().numpy()
evaluator.update_matrix(true, pred)
pbar.update()
if vis_conf_mat:
evaluator.show_matrix(cls_to_idx, save_matrix=True)
return evaluator.error()
