train.back

import torch
import torch.nn as nn
import numpy as np
import math
class BaS_Net_loss(nn.Module):
    def __init__(self, alpha):
        super(BaS_Net_loss, self).__init__()
        self.alpha = alpha
        self.ce_criterion = nn.BCELoss()

    def forward(self, score_base,  fore_weights, label):#score_base:(batch,classes),fore_weights(batch,num_segments,1),label(batch,classes)
        loss = {}

        # print(score_base.size())
        # print(fore_weights.size())
        # print(label.size())
        loss_base = self.ce_criterion(score_base, label)
        # max_att=fore_weights.squeeze().max(dim=-1)[0]
        # min_att=fore_weights.squeeze().min(dim=-1)[0]
        # anomaly_single=label[:,-1]
        # hinge_loss=(1-anomaly_single)*(1-max_att+min_att)+anomaly_single*(max_att-min_att)
        # hinge_loss=hinge_loss.sum()
        
        
        loss_norm = torch.mean(torch.norm(fore_weights, p=1, dim=1))
        loss_total = loss_base  + self.alpha * hinge_loss

        loss["loss_base"] = loss_base
        loss["loss_norm"] = loss_norm
        loss["loss_total"] = loss_total

        return loss_total, loss

def train(net, train_loader, loader_iter, optimizer, criterion, logger, step):
    net.train()
    try:
        _data, _label, _, _, _ = next(loader_iter)
    except:
        loader_iter = iter(train_loader)
        _data, _label, _, _, _ = next(loader_iter)

    _data = _data.cuda()
    _label = _label.cuda()

    optimizer.zero_grad()

    score_base, fore_weights = net(_data)

    cost, loss = criterion(score_base, fore_weights, _label)

    cost.backward()
    optimizer.step()

    for key in loss.keys():
        # logger.log_value(key, loss[key].cpu().item(), step)
        logger.add_scalar(key, loss[key].cpu().item(), step)
        print(step,',',key,":",loss[key].cpu().item())