fix(): change one class softmax implementation, and deactivate Weight…

…edRandomSampling

loss.py

@@ -3,6 +3,117 @@
 from torch.autograd.function import Function
 import torch.nn.functional as F
 from torch.autograd import Variable
+from torch.nn import Parameter
+
+
+class OCAngleLayer(nn.Module):
+    """ Output layer to produce activation for one-class softmax
+
+    Usage example:
+     batchsize = 64
+     input_dim = 10
+     class_num = 2
+
+     l_layer = OCAngleLayer(input_dim)
+     l_loss = OCSoftmaxWithLoss()
+
+     data = torch.rand(batchsize, input_dim, requires_grad=True)
+     target = (torch.rand(batchsize) * class_num).clamp(0, class_num-1)
+     target = target.to(torch.long)
+
+     scores = l_layer(data)
+     loss = l_loss(scores, target)
+
+     loss.backward()
+    """
+
+    def __init__(self, in_planes, w_posi=0.9, w_nega=0.2, alpha=20.0):
+        super(OCAngleLayer, self).__init__()
+        self.in_planes = in_planes
+        self.w_posi = w_posi
+        self.w_nega = w_nega
+        self.out_planes = 1
+
+        self.weight = Parameter(torch.Tensor(in_planes, self.out_planes))
+        # self.weight.data.uniform_(-1, 1).renorm_(2,1,1e-5).mul_(1e5)
+        nn.init.kaiming_uniform_(self.weight, 0.25)
+        self.weight.data.renorm_(2, 1, 1e-5).mul_(1e5)
+
+        self.alpha = alpha
+
+    def forward(self, input, flag_angle_only=False):
+        """
+        Compute oc-softmax activations
+
+        input:
+        ------
+          input tensor (batchsize, input_dim)
+
+        output:
+        -------
+          tuple of tensor ((batchsize, output_dim), (batchsize, output_dim))
+        """
+        # w (feature_dim, output_dim)
+        w = self.weight.renorm(2, 1, 1e-5).mul(1e5)
+        # x_modulus (batchsize)
+        # sum input -> x_modules in shape (batchsize)
+        x_modulus = input.pow(2).sum(1).pow(0.5)
+        # w_modules (output_dim)
+        # w_moduls should be 1, since w has been normalized
+        # w_modulus = w.pow(2).sum(0).pow(0.5)
+
+        # W * x = ||W|| * ||x|| * cos())))))))
+        # inner_wx (batchsize, 1)
+        inner_wx = input.mm(w)
+        # cos_theta (batchsize, output_dim)
+        cos_theta = inner_wx / x_modulus.view(-1, 1)
+        cos_theta = cos_theta.clamp(-1, 1)
+
+        if flag_angle_only:
+            pos_score = cos_theta
+            neg_score = cos_theta
+        else:
+            pos_score = self.alpha * (self.w_posi - cos_theta)
+            neg_score = -1 * self.alpha * (self.w_nega - cos_theta)
+
+        #
+        return pos_score, neg_score
+
+
+class OCSoftmaxWithLoss(nn.Module):
+    """
+    OCSoftmaxWithLoss()
+
+    """
+
+    def __init__(self):
+        super(OCSoftmaxWithLoss, self).__init__()
+        self.m_loss = nn.Softplus()
+
+    def forward(self, inputs, target):
+        """
+        input:
+        ------
+          input: tuple of tensors ((batchsie, out_dim), (batchsie, out_dim))
+                 output from OCAngle
+                 inputs[0]: positive class score
+                 inputs[1]: negative class score
+          target: tensor (batchsize)
+                 tensor of target index
+        output:
+        ------
+          loss: scalar
+        """
+        # Assume target is binary, positive = 1, negaitve = 0
+        #
+        # Equivalent to select the scores using if-elese
+        # if target = 1, use inputs[0]
+        # else, use inputs[1]
+        output = inputs[0] * target.view(-1, 1) + \
+                 inputs[1] * (1 - target.view(-1, 1))
+        loss = self.m_loss(output).mean()
+
+        return loss
 
 
 class OCSoftmax(nn.Module):

main.py

@@ -123,7 +123,7 @@ def pad(x, max_len=64000):
 
 
 def prepare_weights_to_fix_imbalance(dataset, train_ids):
-    class_sample_count = np.array([5000, 45000])
+    class_sample_count = np.array([1, 9])
     weight = 1. / class_sample_count
     samples_weight = []
 
@@ -153,16 +153,16 @@ def train(parser, device):
         model.load_state_dict(torch.load(args.model_path))
         print('Model loaded : {}'.format(args.model_path))
 
-#     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
-#                                  betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)
+    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,
+                                 betas=(args.beta_1, args.beta_2), eps=args.eps, weight_decay=0.0005)
 
-    optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.0005)
+    # optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=0.0005)
 
     train_set = dataset_loader.ASVDataset(is_train=True, transform=transforms)
 
     number_of_epochs = int(args.num_epochs / args.num_folds)
-    checkpoint_epoch = args.epoch % args.num_folds
-    checkpoint_fold = args.epoch // args.num_folds
+    checkpoint_epoch = args.epoch % number_of_epochs
+    checkpoint_fold = args.epoch // number_of_epochs
     monitor_loss = args.add_loss
 
     print(f'{Color.ENDC}Train Start...')
@@ -173,17 +173,17 @@ def train(parser, device):
 
         print(f'{Color.UNDERLINE}{Color.WARNING}Fold {fold}{Color.ENDC}')
         model.train()
-        weights = prepare_weights_to_fix_imbalance(train_set, train_ids)
-        weighted_sampler = torch.utils.data.WeightedRandomSampler(weights=weights, num_samples=len(train_ids),
-                                                                  replacement=True)
-        print(train_set.__getitem__(train_ids[0]))
+        # weights = prepare_weights_to_fix_imbalance(train_set, train_ids)
+        # weighted_sampler = torch.utils.data.WeightedRandomSampler(weights=weights, num_samples=len(train_ids),
+        #                                                           replacement=True)
+
         train_sub_sampler = torch.utils.data.Subset(train_set, train_ids)
         test_sub_sampler = torch.utils.data.Subset(train_set, test_ids)
 
         train_loader = torch.utils.data.DataLoader(
             train_sub_sampler,
-            batch_size=args.batch_size,
-            sampler=weighted_sampler)
+            batch_size=args.batch_size)
+            # sampler=weighted_sampler)
 
         validation_loader = torch.utils.data.DataLoader(
             test_sub_sampler,
@@ -231,7 +231,7 @@ def train(parser, device):
 
                     dev_loss_dict[args.add_loss].append(loss.item())
                     idx_loader.append(labels)
-                    score_loader.append(score)
+                    score_loader.append(score.mean())
 
                 scores = torch.cat(score_loader, 0).data.cpu().numpy()
                 labels = torch.cat(idx_loader, 0).data.cpu().numpy()

model.py

@@ -17,7 +17,31 @@ def __init__(self, input_channels, num_classes, device):
         self.mlp_layer2 = nn.Linear(256, 256).to(device)
         self.mlp_layer3 = nn.Linear(256, 256).to(device)
         self.drop_out = nn.Dropout(0.5)
-        self.oc_softmax = OCSoftmax(256).to(device)
+
+        self.oc_softmax_angle = OCAngleLayer(in_planes=256)
+        self.oc_softmax_loss = OCSoftmaxWithLoss()
+
+        # self.oc_softmax = OCSoftmax(256).to(device)
+
+    def compute_score(self, feature_vec, angle=False):
+        # compute a-softmax output for each feature configuration
+        batch_size = feature_vec.shape[0]
+
+        # negaitve class scores
+        x_cos_val = torch.zeros(
+            [feature_vec.shape[0], 1],
+            dtype=feature_vec.dtype, device=feature_vec.device)
+
+        # positive class scores
+        x_phi_val = torch.zeros_like(x_cos_val)
+
+        s_idx = 0
+        e_idx = batch_size
+        tmp1, tmp2 = self.oc_softmax_angle(feature_vec[s_idx:e_idx], angle)
+        x_cos_val[s_idx:e_idx] = tmp1
+        x_phi_val[s_idx:e_idx] = tmp2
+
+        return [x_cos_val, x_phi_val]
 
     def forward(self, x, labels, is_train=True):
         x = x.to(self.device)
@@ -33,5 +57,8 @@ def forward(self, x, labels, is_train=True):
         self.drop_out(x)
         x = F.relu(self.mlp_layer3(x))
         feat = x
-
-        return self.oc_softmax(feat, labels, is_train)
+
+        scores = self.compute_score(feat, not is_train)
+        loss = self.oc_softmax_loss(scores, labels)
+
+        return loss, scores[0]