🚧 wip

Author: nateraw

tests/base/model_test_steps.py

@@ -115,10 +115,11 @@ def test_step_result_preds(self, batch, batch_idx, optimizer_idx=None):
         result.log('test_loss', loss_test)
         result.log('test_acc', test_acc)
 
-        #lst_of_str = [random.choice(['dog', 'cat']) for i in range(batch_size)]
-        # int_outputs = [random.randint(500, 1000) for i in range(batch_size)]
-        #nested_lst = [[x] for x in int_outputs]
-        #lst_of_dicts = [{k: v} for k, v in zip(lst_of_str, int_outputs)]
+        batch_size = x.size(0)
+        lst_of_str = [random.choice(['dog', 'cat']) for i in range(batch_size)]
+        lst_of_int = [random.randint(500, 1000) for i in range(batch_size)]
+        lst_of_lst = [[x] for x in lst_of_int]
+        lst_of_dict = [{k: v} for k, v in zip(lst_of_str, lst_of_int)]
 
         # This is passed in from pytest via parameterization
         option = getattr(self, 'test_option', 0)
@@ -135,5 +136,31 @@ def test_step_result_preds(self, batch, batch_idx, optimizer_idx=None):
         elif option == 1:
             result.write('idxs', torch.cat((lazy_ids, lazy_ids)), prediction_file)
             result.write('preds', labels_hat, prediction_file)
+        
+        # write multi-dimension
+        elif option == 2:
+            result.write('idxs', lazy_ids, prediction_file)
+            result.write('preds', labels_hat, prediction_file)
+            result.write('x', x, prediction_file)
+        
+        # write str list
+        elif option == 3:
+            result.write('idxs', lazy_ids, prediction_file)
+            result.write('vals', lst_of_str, prediction_file)
+
+        # write int list
+        elif option == 4:
+            result.write('idxs', lazy_ids, prediction_file)
+            result.write('vals', lst_of_str, prediction_file)
+
+        # write nested list
+        elif option == 5:
+            result.write('idxs', lazy_ids, prediction_file)
+            result.write('vals', lst_of_str, prediction_file)
+
+        # write dict list
+        elif option == 6:
+            result.write('idxs', lazy_ids, prediction_file)
+            result.write('vals', lst_of_str, prediction_file)
 
         return result

tests/core/test_results.py

@@ -55,6 +55,21 @@ def test_result_reduce_ddp(result_cls):
         pytest.param(
             1, False, 0, id='test_only_mismatching_tensor', marks=pytest.mark.xfail(raises=ValueError, match="Mism.*")
         ),
+        pytest.param(
+            2, False, 0, id='mix_of_tensor_dims'
+        ),
+        pytest.param(
+            3, False, 0, id='string_list_predictions'
+        ),
+        pytest.param(
+            4, False, 0, id='int_list_predictions'
+        ),
+        pytest.param(
+            5, False, 0, id='nested_list_predictions'
+        ),
+        pytest.param(
+            6, False, 0, id='dict_list_predictions'
+        ),
         pytest.param(
             0, True, 1, id='full_loop_single_gpu', marks=pytest.mark.skipif(torch.cuda.device_count() < 1, reason="test requires single-GPU machine")
         )
@@ -153,7 +168,9 @@ def test_result_obj_predictions_ddp_spawn(tmpdir):
     dm.setup('test')
 
     # check prediction file now exists and is of expected length
+    size = 0
     for prediction_file in prediction_files:
         assert prediction_file.exists()
         predictions = torch.load(prediction_file)
-        assert len(predictions) == len(dm.mnist_test) // 2
+        size += len(predictions)
+    assert size == len(dm.mnist_test)