├── README.md
└── TSA_crossentropy_loss.py


/README.md:
--------------------------------------------------------------------------------
 1 | # TSA_pytorch
 2 | Training Signal Annealing
 3 | 
 4 | This is an example of TSA.
 5 | In the code above, I compare training a simple neural net of binary classification task:  
 6 | *Input:*  
 7 | ```[a,b]```  
 8 | *Output:*  
 9 | ```1 if a>b```  
10 | ```0 if a<b```  
11 |   
12 | ## Training data  
13 | ```X_train = [11,10],[200,151],[501,400],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12],[12,13],[1,10],[9,20],[9,100],[100,101],[1111,1112]]```  
14 | ```Y_train = [1,1,1,0,0,0,0,0,0,0,0,0,0,0,0]```  
15 | 
16 | ## Result  
17 | ```X_test = [[200,151],[501,400],[9,8],[99,95],[100,100000],[99,60]]```  
18 | ```Y_test = [1,1,1,1,0,1]```  
19 | TSA_training:   
20 | ```output: [[0., 1.],[0., 1.],[0., 1.],[0., 1.],[1., 0.],[0., 1.]]```  
21 |   
22 | Without TSA_training:  
23 | ```output: [[0., 1.],[0., 1.],[0., 1.],[1., 0.],[1., 0.],[0., 1.]]```  
24 |   
25 | 
26 | In the example, the number of "1"-class samples is much smaller than that of "0"-class samples.  
27 | The simple test show that, the Net dose not recognize [99,95] input when it is trained in normal way while the same Net (with TSA training) can classify that input.  
28 | 
29 | ## Conclution  
30 | The code is a demo of applying TSA into crossEntopyLoss written in Pytorch.  
31 | TSA make *loss surface* becomes *smoother* which means of *fewer local minimums*, helping the process of training neural net become easier, expecially in case of difficulty in data collection.  
32 | The experiment implies that Neural Net can be trained with *unbalanced* dataset.   
33 | 
34 | ## Reference  
35 | Qizhe Xie, Zihang Dai, Eduard Hovy, Minh-Thang Luong, Quoc V. Le. Unsupervised Data Augmentation for Consistency Training https://arxiv.org/abs/1904.12848  
36 | 


--------------------------------------------------------------------------------
/TSA_crossentropy_loss.py:
--------------------------------------------------------------------------------
  1 | import torch
  2 | import torch.nn.functional as fn
  3 | 
  4 | 
  5 | class TSA_crossEntropy(object):
  6 |     def __init__(self, num_steps, num_classes, alpha, temperature, cuda = False):
  7 |         self.loss_function = torch.nn.CrossEntropyLoss(reduction='none')
  8 |         self.num_steps = num_steps
  9 |         self.current_step = 0
 10 |         self.num_classes = num_classes
 11 |         self.alpha = alpha
 12 |         self.cuda = cuda
 13 |         self.temperature = temperature
 14 |     
 15 |     def thresh(self):
 16 |         thresh = torch.exp(torch.Tensor([(self.current_step/self.num_steps-1)*self.alpha]))*(1-1/self.num_classes)+1/self.num_classes
 17 |         if self.cuda:
 18 |             thresh = thresh.cuda()
 19 |         return thresh
 20 | 
 21 |     def step(self):
 22 |         self.current_step += 1
 23 | 
 24 |     def get_mask(self, logits, targets):
 25 |         thresh = self.thresh()
 26 |         mask = fn.softmax(logits, dim = 1).detach()
 27 |         mask, pred = torch.max(mask, dim = 1, keepdim = False)
 28 | 
 29 |         wrong_pred = (torch.abs(pred - targets)>0)
 30 |         mask = (mask<thresh)
 31 |         mask |= wrong_pred
 32 | 
 33 |         if self.cuda:
 34 |             return mask.type(torch.cuda.FloatTensor)
 35 |         else:
 36 |             return mask.type(torch.FloatTensor)
 37 | 
 38 | 
 39 |     def __call__(self, logits, targets):
 40 |         logits = logits/self.temperature
 41 | 
 42 |         mask = self.get_mask(logits, targets)
 43 |         self.loss_value = self.loss_function(logits, targets)*mask
 44 |         self.loss_value = self.loss_value.sum()/torch.max(torch.Tensor([mask.sum(),1]))
 45 |         self.step()
 46 |         return self.loss_value
 47 | 
 48 | 
 49 | # Test TSA with crossentropy loss
 50 | loss = TSA_crossEntropy(10000,2,5, 5)
 51 | 
 52 | inputs = torch.FloatTensor([[11,10],[200,151],[501,400],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12],[12,13],[1,10],[9,20],[9,100],[100,101],[1111,1112]])
 53 | weight = torch.FloatTensor([[0,0],[0,0]]).requires_grad_(requires_grad=True)
 54 | bias = torch.FloatTensor([[0,0]]).requires_grad_(requires_grad=True)
 55 | targets = torch.LongTensor([1,1,1,0,0,0,0,0,0,0,0,0,0,0,0])
 56 | 
 57 | 
 58 | for i in range(10000):
 59 |     logits = torch.matmul(inputs, weight)
 60 |     logits = logits+bias
 61 |     loss_value = loss(logits, targets)
 62 | 
 63 |     loss_value.backward()
 64 | 
 65 |     optimizer = torch.optim.SGD([weight, bias], lr=0.01, momentum=0.9)
 66 | 
 67 |     optimizer.step()
 68 | 
 69 | inputs = torch.FloatTensor([[200,151],[501,400],[9,8],[99,95],[100,100000],[99,60]])
 70 | 
 71 | logits = torch.matmul(inputs, weight)
 72 | logits = torch.nn.functional.softmax(logits, dim = 1)
 73 | out,_ = torch.max(logits, dim =1)
 74 | print('with TSA training')
 75 | print('weights:', weight)
 76 | print('bias:', bias)
 77 | print('output:', logits)
 78 | 
 79 | 
 80 | # without TSA 
 81 | 
 82 | loss = torch.nn.CrossEntropyLoss()
 83 | 
 84 | inputs = torch.FloatTensor([[11,10],[200,151],[501,400],[6,7],[7,8],[8,9],[9,10],[10,11],[11,12],[12,13],[1,10],[9,20],[9,100],[100,101],[1111,1112]])
 85 | weight = torch.FloatTensor([[0,0],[0,0]]).requires_grad_(requires_grad=True)
 86 | bias = torch.FloatTensor([[0,0]]).requires_grad_(requires_grad=True)
 87 | targets = torch.LongTensor([1,1,1,0,0,0,0,0,0,0,0,0,0,0,0])
 88 | 
 89 | 
 90 | for i in range(10000):
 91 |     logits = torch.matmul(inputs, weight)
 92 |     logits = logits+bias
 93 |     loss_value = loss(logits, targets)
 94 | 
 95 |     loss_value.backward()
 96 | 
 97 |     optimizer = torch.optim.SGD([weight, bias], lr=0.01, momentum=0.9)
 98 | 
 99 |     optimizer.step()
100 | 
101 | inputs = torch.FloatTensor([[200,151],[501,400],[9,8],[99,95],[100,100000],[99,60]])
102 | logits = torch.matmul(inputs, weight)
103 | logits = torch.nn.functional.softmax(logits, dim = 1)
104 | out,_ = torch.max(logits, dim =1)
105 | print('without TSA training')
106 | print('weights:', weight)
107 | print('bias:', bias)
108 | print('output:', logits)
109 | 
110 | 
111 | 


--------------------------------------------------------------------------------