Training the network with some batch size - code

There is my "training" code below, I wrote it based on one youtube tutorial. I don't understand actually one part: batch_X = train_X[i:i+BATCH_SIZE], batch_y = train_y[i:i+BATCH_SIZE]. How to understand this thing inside the brackets (i:i+BATCH_SIZE)?

def train(net):
BATCH_SIZE = 32
EPOCHS = 10
for epoch in range(EPOCHS):
# training loop
net.train()
for i in tqdm(range(0, len(train_X), BATCH_SIZE)):
batch_X = train_X[i:i+BATCH_SIZE].view(-1,3,224,224)
batch_y = train_y[i:i+BATCH_SIZE]

batch_X, batch_y = batch_X.to(device), batch_y.to(device)

net.zero_grad()
outputs1 = net(batch_X)
matches1 = [torch.argmax(i)==torch.argmax(j) for i, j in zip(outputs1, batch_y)]
acc = matches1.count(True)/len(matches1)
loss = loss_function(outputs1, batch_y)
loss.backward()
optimizer.step()

print(f"Epoch: {epoch}. Training acc: {round(float(acc),2)} Training loss: {round(float(loss),3)}")


1 Answer

That part of the code will select the samples that belong to a specific batch. The for loop first loops over the data in train_X in steps of BATCH_SIZE, which means that the variable i holds the first index for each batch in the training dataset. The rest of the samples for the batch are then the ones after that index up to the sample which completes the batch. This is done using train_X[i:i+BATCH_SIZE].