Is it cheating to stratified split the whole dataset based on a previous test set result?

Question

I trained a model using a small mri dataset(57 patients). The model's performance was so low(Train set 0.7, Val set 0.7, Test set 0.45).

I found the model segments tumor in upper part of brain well, couldn't segment tumor which is in middle part of brain(information from validation + test).

So i stratified split the whole dataset with a position stratum(upper, not upper), the model performance was improved. I test this 3 times with different random states for split, average test accuracy increased a little.

In this case, is it cheating to do that?

Answer 1

First of all, the first scores train set 0.7, val set 0.7, test set 0.45 shows that your model suffers from overfitting.

Because of you have a really small dataset you can try one leave out cross-validation technique to train your model.

Another thing that you need to be careful is about the distribution of your target. From your question, I understand that your problem is binary but what is the 0/1 ratio?

I can edit my answer based on this.

Ex: 57 samples, and the distribution of your all dataset is like below.

20 examples from the upper part of the brain.
25 examples from the middle part of the brain.
12 examples from the lower part of the brain.

In this case, you can split the train/test (80:20) like below to ensure the same train/test distribution.

Train set: Total Example: 46 (57 * 80/100)
46 example contains: 
1. 20 * 80/100 = 16 upper part sample
2. 25 * 80/100 = 20 middle part sample
3. 12 * 80/100 = 10 lower part sample

Testset: Total Example: 11 (57 * 20/100)
1. 20 * 20/100 = 4 upper part sample
2. 25 * 20/100 = 5 middle part sample
3. 12 * 20/100 = 2 lower part sample