Some problems with WeightedRandomSampler

yu_Magic: <p>Dear groupers,<br>
I work on an unbalanced dataset. There are six class in my dataset. The first class has 568330 samples, the second class has 43000 samples, the third class has 34900, the fourth class has 20910, the fifth class has 14590, and the last class has 9712 class. I used WeighedRandomSampler in my dataloader. The purpose of my dataloader is each class can sampling averagely.  But i found when i used WeightedRandomSampler, the class in one batch size always is class first class, second class and third class. No fourth and fifth class in each batchsize. Here is my code snippet.<br>
class_sample_counts = [568330.0, 43000.0, 34900.0, 20910.0, 14590.0, 9712.0]<br>
class_weights = 1./torch.tensor(class_sample_counts)<br>
sampler = torch.utils.data.sampler.WeightedRandomSampler(class_weights, num_samples=len(my_dataset), replacement=True)<br>
loader = torch.utils.data.DataLoader(<br>
dataset=my_dataset,<br>
batch_size=batch_size,<br>
sampler=sampler,<br>
pin_memory=False,<br>
num_workers=number_workers,<br>
)<br>
Can anyone help me to check my problems? Thanks a lot!</p>

ptrblck: <p>I think you might pass the wrong weights to <code>WeightedRandomSampler</code>.<br>
The sequence of weights should correspond to your samples in the dataset.<br>
Here is a small example:</p>
<pre><code class="lang-python">weights = 1. / torch.tensor(class_sample_counts, dtype=torch.float)
samples_weights = weights[train_targets]

sampler = WeightedRandomSampler(
    weights=samples_weights,
    num_samples=len(samples_weights),
    replacement=True)
</code></pre>

yu_Magic: <p>Thank you for your prompt reply and your kindest help! I still confuse about your small example. Could you please explain more details about the train_targets in the sample_weights = weights[train_targets]?</p>

ptrblck: <p>This line of code creates a tensor containing the corresponding weight value for each sample.<br>
While <code>weights</code> has the shape <code>[num_classes]</code>, <code>sample_weights</code> has <code>[num_samples]</code> containing <code>[weights[y0], weights[y1], ...]</code>.</p>

surojit_sengupta: <p>Hello Peter,<br>
(I hope that’s your intended name)</p>
<p>Am doing something similar but, my target isn’t just a class it’s rather a tensor of classes.</p>
<p>data1 = ([‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’, ‘boycott’, ‘British’, ‘lamb’, ‘.’],[‘B-ORG\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’])<br>
The first list is the sentence while the second is the list of classes for each word.</p>
<p>This is further passed into a windower to give an output like so: (with relevant tags)</p>
<p>[’&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’]<br>
[’&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’, ‘boycott’]<br>
[’&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’, ‘boycott’, ‘British’]<br>
[’&lt;PAD&gt;’, ‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’, ‘boycott’, ‘British’, ‘lamb’]<br>
[‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’, ‘boycott’, ‘British’, ‘lamb’, ‘.’]</p>
<p>Now sampling is messing up with my sentence structure in a very wrong way and hence the word contexts for an NER task are lost.</p>
<p>Check an example below:</p>
<p>[(’&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘German’, ‘EU’, ‘&lt;PAD&gt;’, ‘German’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘rejects’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘to’, ‘German’, ‘rejects’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’), (’&lt;PAD&gt;’, ‘EU’, ‘EU’, ‘call’, ‘rejects’, ‘&lt;PAD&gt;’, ‘call’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘EU’, ‘&lt;PAD&gt;’, ‘EU’, ‘German’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘boycott’, ‘call’, ‘German’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’), (‘EU’, ‘rejects’, ‘rejects’, ‘to’, ‘German’, ‘&lt;PAD&gt;’, ‘to’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘rejects’, ‘EU’, ‘rejects’, ‘call’, ‘EU’, ‘EU’, ‘British’, ‘to’, ‘call’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’), (‘rejects’, ‘German’, ‘German’, ‘boycott’, ‘call’, ‘&lt;PAD&gt;’, ‘boycott’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘German’, ‘rejects’, ‘German’, ‘to’, ‘rejects’, ‘rejects’, ‘lamb’, ‘boycott’, ‘to’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’), (‘German’, ‘call’, ‘call’, ‘British’, ‘to’, ‘Peter’, ‘British’, ‘Peter’, ‘Peter’, ‘call’, ‘German’, ‘call’, ‘boycott’, ‘German’, ‘German’, ‘.’, ‘British’, ‘boycott’, ‘Peter’, ‘Peter’), (‘call’, ‘to’, ‘to’, ‘lamb’, ‘boycott’, ‘Blackburn’, ‘lamb’, ‘Blackburn’, ‘Blackburn’, ‘to’, ‘call’, ‘to’, ‘British’, ‘call’, ‘call’, ‘&lt;PAD&gt;’, ‘lamb’, ‘British’, ‘Blackburn’, ‘Blackburn’), (‘to’, ‘boycott’, ‘boycott’, ‘.’, ‘British’, ‘&lt;PAD&gt;’, ‘.’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘boycott’, ‘to’, ‘boycott’, ‘lamb’, ‘to’, ‘to’, ‘&lt;PAD&gt;’, ‘.’, ‘lamb’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’), (‘boycott’, ‘British’, ‘British’, ‘&lt;PAD&gt;’, ‘lamb’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘British’, ‘boycott’, ‘British’, ‘.’, ‘boycott’, ‘boycott’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘.’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’), (‘British’, ‘lamb’, ‘lamb’, ‘&lt;PAD&gt;’, ‘.’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘lamb’, ‘British’, ‘lamb’, ‘&lt;PAD&gt;’, ‘British’, ‘British’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’, ‘&lt;PAD&gt;’)] [(‘PAD’, ‘PAD’, ‘PAD’, ‘B-MISC\n’, ‘B-ORG\n’, ‘PAD’, ‘B-MISC\n’, ‘PAD’, ‘PAD’, ‘PAD’, ‘PAD’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘PAD’, ‘PAD’), (‘PAD’, ‘B-ORG\n’, ‘B-ORG\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘B-ORG\n’, ‘PAD’, ‘B-ORG\n’, ‘B-MISC\n’, ‘PAD’, ‘PAD’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘PAD’, ‘PAD’), (‘B-ORG\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘O\n’, ‘B-ORG\n’, ‘O\n’, ‘O\n’, ‘B-ORG\n’, ‘B-ORG\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘PAD’), (‘O\n’, ‘B-MISC\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘B-MISC\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘PAD’), (‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘B-PER\n’, ‘B-MISC\n’, ‘B-PER\n’, ‘B-PER\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘B-MISC\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘B-PER\n’, ‘B-PER\n’), (‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘I-PER\n’, ‘O\n’, ‘I-PER\n’, ‘I-PER\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘O\n’, ‘B-MISC\n’, ‘I-PER\n’, ‘I-PER\n’), (‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘O\n’, ‘O\n’, ‘PAD’, ‘PAD’), (‘O\n’, ‘B-MISC\n’, ‘B-MISC\n’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘PAD’, ‘PAD’, ‘B-MISC\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’), (‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘PAD’, ‘O\n’, ‘PAD’, ‘PAD’, ‘PAD’, ‘PAD’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘PAD’, ‘B-MISC\n’, ‘B-MISC\n’, ‘PAD’, ‘PAD’, ‘PAD’, ‘PAD’, ‘PAD’)</p>
<p>Regards</p>

ptrblck: <p>Could you provide the <code>Dataset</code> code and how you are calculating the weights?<br>
I’m not an NLP expert, but I guess your sentences shouldn’t be split by the sampling, so it’s interesting to see, how you are dealing with your data generally.<br>
Also, you might want to have a look at <a href="https://github.com/pytorch/text/tree/master/test">torchtext</a>.</p>

surojit_sengupta: <p>Hello Peter,</p>
<p>This is a conll2003 corpus and a sentence looks like one below:</p>
<p>EU NNP B-NP B-ORG<br>
rejects VBZ B-VP O<br>
German JJ B-NP B-MISC<br>
call NN I-NP O<br>
to TO B-VP O<br>
boycott VB I-VP O<br>
British JJ B-NP B-MISC<br>
lamb NN I-NP O</p>
<p>All I have to do is extract the word and it’s NER tag in the right extreme and make my sentence look like one below:</p>
<p>([‘EU’, ‘rejects’, ‘German’, ‘call’, ‘to’, ‘boycott’, ‘British’, ‘lamb’, ‘.’] → One Image<br>
[‘B-ORG\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’, ‘O\n’, ‘B-MISC\n’, ‘O\n’, ‘O\n’]) → set of tags/labels for an image</p>
<p>Each unique word and tag has an index like so:</p>
<blockquote>
<p>{‘B-LOC\n’: 4,<br>
‘B-MISC\n’: 9,<br>
‘B-ORG\n’: 6,<br>
‘B-PER\n’: 7,<br>
‘I-LOC\n’: 8,<br>
‘I-MISC\n’: 3,<br>
‘I-ORG\n’: 1,<br>
‘I-PER\n’: 5,<br>
‘O\n’: 2,<br>
‘PAD’: 0}</p>
</blockquote>
<p>Here each word in a sentence is converted to its indices like so:</p>
<blockquote>
<p>[    0, 10397, 21273, 16075, 20124, 15819,  2729, 13298, 14387]</p>
</blockquote>
<p>; and is embedded through a nn.Embedding layer (so that each word is a 50 dimensional vector)</p>
<p>and now a single sentence becomes a 2D image. But, instead of having a single label to an image, I have a list of labels to an image. These labels are the NER tags of each word. The problem is, my data-set has a lot of words of  ‘O\n’ class as pointed in the comment earlier and so, my model tends to predict the dominant class (typical class imbalance problem).<br>
So, I need to balance these classes.</p>
<p>The code to calculate weights:</p>
<p><code>indexed_counts  #frequency of each class</code></p>
<blockquote>
<p>{0: 112328,<br>
1: 3704,<br>
2: 169578,<br>
3: 1155,<br>
4: 7140,<br>
5: 4528,<br>
6: 6321,<br>
7: 6600,<br>
8: 1157,<br>
9: 3438}</p>
</blockquote>
<pre><code class="lang-auto">tag_weights = {}
for key in indexed_counts.keys():
    tag_weights[key] = 1/indexed_counts[key]
sampler = [i[1] for i in sorted(tag_weights.items())]
sampler
</code></pre>
<blockquote>
<p>[8.902499821950003e-06,<br>
0.0002699784017278618,<br>
5.896991355010673e-06,<br>
0.0008658008658008658,<br>
0.00014005602240896358,<br>
0.00022084805653710247,<br>
0.00015820281601012498,<br>
0.00015151515151515152,<br>
0.000864304235090752,<br>
0.00029086678301337986]</p>
</blockquote>
<p>(As a side note, am passing this as my weights tensor for the cross entropy loss function)</p>
<p>Just for you to compare that the tag_weights and the sampler is in the same order of the weight classes<br>
<code>tag_weights</code></p>
<blockquote>
<p>{<strong>0</strong>: 8.902499821950003e-06, <strong>1</strong>: 0.0002699784017278618, <strong>2</strong>: 5.896991355010673e-06, <strong>3</strong>: 0.0008658008658008658, <strong>4</strong>: 0.00014005602240896358, <strong>5</strong>: 0.00022084805653710247, <strong>6</strong>: 0.00015820281601012498, <strong>7</strong>: 0.00015151515151515152, <strong>8</strong>: 0.000864304235090752, <strong>9</strong>: 0.00029086678301337986}</p>
</blockquote>
<p>Coming to your suggestion of splitting by classes, splitting a sentence without maintaining the sequence of words before and after, would imply the context of a word might be lost. The model needs to learn the NER tags of a given word given the contextual meaning as well (words appearing before and after a given word).</p>

surojit_sengupta: <p>On a different note, is there a way to penalize my model basing upon the F1 score? Instead of considering sampling as a solution, can we somehow train our model towards a better F1 score? So the point being, how to choose the f1 score as a metric for my model instead of the loss only?</p>

sidd2006: <p>You can keep a track of F1 on the validation data and decide when to stop the training based on the best F1. I am using ignite library to do a similar thing. You can create a metric: F1Weighted(Metric): using the Metric class from <a href="https://pytorch.org/ignite/" rel="nofollow noopener">pytorch-ignite</a>. Then, add the metric as follows:</p>
<pre><code class="lang-auto">evaluator = create_supervised_evaluator(model,
                                                metrics={'f_1': F1Weighted})
</code></pre>
<p>You can use earlystopping from ignite to stop training and also modelcheckpointing. Its all present in ignite.</p>

surojit_sengupta: <p>Thanks for the new lead on a new day! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=6" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
I’ll get back to you for nay help, if needed,soon.</p>

Wafaa_Wardah: <p>I have unbalanced data and want to use this solution. However, I get this error:</p>
<p>UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).<br>
self.weights = torch.tensor(weights, dtype=torch.double)</p>
<p>Can you please help?</p>
<pre><code class="lang-auto">classcount = np.bincount(trainset.labels).tolist()
train_weights = 1./torch.tensor(classcount, dtype=torch.float)
train__sampleweights = train_weights[trainset.labels]
train_sampler = Sampler.WeightedRandomSampler(weights=train__sampleweights, num_samples=len(train__sampleweights))
trainloader = DataLoader(trainset, sampler=train_sampler, shuffle=False)
</code></pre>

ptrblck: <p>It looks like <code>train_sampleweights</code> is a <code>numpy</code> array.<br>
Could you convert it to a <code>torch.tensor</code> using <code>train_sampleweights = torch.from_numpy(train_sampleweights)</code> before passing it to the <code>WeightedRandomSampler</code>?</p>

Wafaa_Wardah: <aside class="quote no-group" data-username="ptrblck" data-post="12" data-topic="23242">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>train_sampleweights = torch.from_numpy(train_sampleweights)</p>
</blockquote>
</aside>
<p>Thank you for replying. Unfortunately, it gave this error:</p>
<pre><code class="lang-auto">train_sampleweights = torch.from_numpy(train_sampleweights)
TypeError: expected np.ndarray (got Tensor)
</code></pre>
<p>So when I print out some information on train_sampleweights, I think it may help, but I’m not sure what the class constructor wants.</p>
<pre><code class="lang-auto">train_weights = tensor([4.3002e-06, 4.3002e-06, 4.3002e-06,  ..., 4.3002e-06, 4.3002e-06, 4.3002e-06])
train_weights shape=  torch.Size([246048])
</code></pre>

ptrblck: <p>Thanks for the info. I was mistaken and your code should work fine.<br>
You can ignore this warning for now, as it’s being thrown internally in <a href="https://github.com/pytorch/pytorch/blob/279238f0b81dce43636ebc8c7b1c45cc98c7a60d/torch/utils/data/sampler.py#L119" rel="nofollow noopener">this line of code</a>.<br>
I don’t think it will cause any trouble.</p>
<p>CC <a class="mention" href="/u/simonw">@SimonW</a> I think this change was made in <a href="https://github.com/pytorch/pytorch/pull/6249/" rel="nofollow noopener">this PR</a>.<br>
Do you think it should be changed?</p>

SimonW: <p>oh sure. the warning was introduced later after that PR was merged, but we should change this. Could you open an issue?</p>

ptrblck: <p>Sure!<br>
I’ll open an issue and suggest a fix.</p>

Alessio_Rosatelli: <p>Hello everyone,</p>
<p>I’m replying here since my problem is like the first question asked and I avoid opening another thread.</p>
<p>I’m struggling with a 3 classes problem with great unbalance. In the following code I create the weight array ( first the array with dimension [ num_classes] then the one with dimension [data_length]</p>
<pre><code class="lang-auto">weight = np.zeros(self.num_classes)
weight = 1. / count_train
self.train_samples_weight = [weight[cla] for cla in Labels_train]
self.test_samples_weight = [weight[class_id] for class_id in Labels_test]
self.train_samples_weight = np.asarray(self.train_samples_weight)
self.test_samples_weight = np.asarray(self.test_samples_weight)
</code></pre>
<p>and it looks correct to me. Then  the code were I create the Sampler and the Loader</p>
<pre><code class="lang-auto">train_sampler = torch.utils.data.WeightedRandomSampler(self.train_samples_weight,1,replacement=True)
test_sampler = torch.utils.data.WeightedRandomSampler(self.test_samples_weight,1,replacement=True)
test_loader = DataLoader(dataset=self.dataset_test, batch_size=self.batch_size, sampler=test_sampler,drop_last=True)
train_loader = DataLoader(dataset=self.dataset_train, batch_size=self.batch_size, sampler=train_sampler,drop_last=True)
</code></pre>
<p>The first thing is I’m using 1 as num_samples since I want just extract a number of samples correspondant to my batch size. Anyway this is missleading since the other Sampler doesn’t work like this as pointed out <a href="https://github.com/pytorch/pytorch/issues/9171" rel="nofollow noopener">here</a> but I guess the problem is in my pytorch version.</p>
<p>The real issue resides in the class distribution when i use enumerate on my DataLoader. In fact I can only get samples of class_2 ( that’s the class with major number of occurrencies) but its weight seems to be the lowest and that looks correct to me. So what I’m doing wrong? I was looking for a way to express a balanced batch size and I still think that’s the correct way to go.</p>

ptrblck: <p><code>num_samples</code> gives the number of samples to draw, so usually you would let it to the length of your <code>Dataset</code>. Currently you should only get a single batch containing one sample.</p>

Alessio_Rosatelli: <p>Yes, I just checked this and that’s True. Anyway I can’t understand why num_samples should be equal to my dataset dimension. In fact I already viewed example with this settings and that make no sense to me.</p>
<p>Without specifing the Sampler my enumerate(loader) results in batchID,data where data has dimension equal to [batchSize,…My data dimension…] besides, when I use customized Sampler I end up having a data dimneison of [ batchSize, numSamples,…My data dimension…].</p>
<p>Can you explain why this should work like that? I can’t see the point in this. I was expecting to still have same dimensions but with samples choosed according to new criteria.</p>

ptrblck: <p>That shouldn’t be the case and I’ve never seen this behavior before.<br>
Could you post a code snippet, so that we can reproduce this issue?<br>
Here is a dummy example which results in <code>[batch_size, nb_features]</code> for each batch:</p>
<pre><code class="lang-python"># Create dummy data with class imbalance 99 to 1
numDataPoints = 1000
data_dim = 5
bs = 100
data = torch.randn(numDataPoints, data_dim)
target = torch.cat((torch.zeros(int(numDataPoints * 0.99), dtype=torch.long),
                    torch.ones(int(numDataPoints * 0.01), dtype=torch.long)))

print('target train 0/1: {}/{}'.format(
    (target == 0).sum(), (target == 1).sum()))

# Compute samples weight (each sample should get its own weight)
class_sample_count = torch.tensor(
    [(target == t).sum() for t in torch.unique(target, sorted=True)])
weight = 1. / class_sample_count.float()
samples_weight = torch.tensor([weight[t] for t in target])

# Create sampler, dataset, loader
sampler = WeightedRandomSampler(samples_weight, len(samples_weight))
train_dataset = torch.utils.data.TensorDataset(data, target)
train_loader = DataLoader(
    train_dataset, batch_size=bs, num_workers=1, sampler=sampler)

# Iterate DataLoader and check class balance for each batch
for i, (x, y) in enumerate(train_loader):
    print("batch index {}, 0/1: {}/{}".format(
        i, (y == 0).sum(), (y == 1).sum()))
    print("x.shape {}, y.shape {}".format(x.shape, y.shape))
</code></pre>