switch to sparse one hot label

Author: bindog

.gitignore

@@ -1,3 +1,6 @@
+# Other temp files
+*.swp
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

cross_entropy.py

@@ -53,15 +53,10 @@ def forward(ctx, *args):
         if ctx.compute_loss:
             _loss_list = []
             for gpu_id, softmax in enumerate(softmax_list):
-                if isinstance(ctx.label_split[gpu_id], torch.sparse.LongTensor):
-                    idx = ctx.label_split[gpu_id]._indices()
-                    # FIXME move _loss to gpu?
-                    _loss = torch.zeros(ctx.batch_size)
-                    _loss.scatter_(dim=0, index=idx[0], src=softmax[tuple(idx)])
-                    _loss_list.append(_loss)
-                else:
-                    _loss = torch.sum(softmax * ctx.label_split[gpu_id], dim=1)
-                    _loss_list.append(_loss)
+                idx = ctx.label_split[gpu_id]._indices()
+                _loss = torch.zeros(ctx.batch_size).to(gpu_id)
+                _loss.scatter_(dim=0, index=idx[0], src=softmax[tuple(idx)])
+                _loss_list.append(_loss)
             _loss = comm.reduce_add(_loss_list, destination=0)
             log_loss = -torch.log(_loss)
             loss = torch.mean(log_loss)

train.py

@@ -14,7 +14,7 @@
 
 from model import ft_net
 from cross_entropy import ModelParallelCrossEntropy
-from utils import get_class_split, get_onehot_label, compute_batch_acc
+from utils import get_class_split, get_sparse_onehot_label, compute_batch_acc
 
 
 def get_data_loader(data_path, batch_size):
@@ -45,7 +45,7 @@ def train_model(opt, data_loader, model, criterion, optimizer, class_split):
             images, labels = data_loader_iter.next()
             images = images.cuda(0)
             labels = labels.cuda(0)
-            onehot_labels = get_onehot_label(labels, opt.num_gpus, opt.num_classes, opt.model_parallel, class_split)
+            onehot_labels = get_sparse_onehot_label(labels, opt.num_gpus, opt.num_classes, opt.model_parallel, class_split)
             # Forward
             optimizer.zero_grad()
             logits = model(images, labels=onehot_labels)