def train_fn(model: torch.nn, data_loader: DataLoader, optimizer: optim,
device: torch.device, epoch: int):
model.train()
start_time = datetime.datetime.now()
num_images: int = 0
for i, (images, targets) in enumerate(data_loader):
images = list(image.to(device) for image in images)
targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
images = torch.stack(images)
num_images += len(images)
optimizer.zero_grad()
loss_dict: Dict[str, torch.Tensor] = model(images, targets)
loss: float = sum(loss for loss in loss_dict.values())
loss.backward()
optimizer.step()
if (i + 1) % 10 == 0:
print('-' * 50)
print(
f'Epoch {epoch+1}[{len(data_loader.dataset):,}/{(num_images/len(data_loader.dataset))*100:.2f}%] '
f'- Elapsed time: {datetime.datetime.now() - start_time}\n'
f' - loss: classifier={loss_dict["loss_classifier"]:.6f}, box_reg={loss_dict["loss_box_reg"]:.6f}, '
f'objectness={loss_dict["loss_objectness"]:.6f}, rpn_box_reg={loss_dict["loss_rpn_box_reg"]:.6f}'
)
def _get_cv_stats(self, model: torch.nn) -> Tuple[torch.tensor, ...]:
"""
Collect CV data accuracy, loss, prediction distribution, and targets
distribution.
Arguments:
model {torch.nn} -- model to learn
Returns:
Tuple[torch.tensor * 4] -- accuracy, loss, prediction distribution, and
targets distribution.
"""
cv_acc = torch.tensor(0.0).to(self.device)
cv_loss = torch.tensor(0.0).to(self.device)
samples_no = float(len(self.cvloader.dataset))
outputs_dist = None
targets_dist = None
with torch.no_grad():
model = model.eval()
for inputs, targets in self.cvloader:
batch_size = inputs.shape[
0] # last sample can have different items
targets = targets.to(self.device)
outputs = model(inputs.to(self.device))
if outputs_dist is None and targets_dist is None:
outputs_dist = outputs.argmax(1).long()
targets_dist = targets.long()
else:
outputs_dist = torch.cat(
[outputs_dist, outputs.argmax(1).long()])
targets_dist = torch.cat([targets_dist, targets.long()])
cv_acc += (outputs.argmax(1) == targets).sum()
cv_loss += self.criterion(outputs, targets) * batch_size
cv_acc = (cv_acc / samples_no).to(CPU_DEVICE)
cv_loss = (cv_loss / samples_no).to(CPU_DEVICE)
outputs_dist = outputs_dist.to(CPU_DEVICE)
targets_dist = targets_dist.to(CPU_DEVICE)
model = model.train()
return cv_acc, cv_loss, outputs_dist, targets_dist