Stack Exchange
How to add one more batch size dimension in python class code
I am learning attention mechanism and saw someone's code from github as below:
class PositionAttentionModule(nn.Module):
    def __init__(self, in_channels):
        super(PositionAttentionModule, self).__init__()
        self.first_branch_conv = nn.Conv2d(in_channels, int(in_channels/8), kernel_size = 1)
        self.second_branch_conv = nn.Conv2d(in_channels, int(in_channels/8), kernel_size = 1)
        self.third_branch_conv = nn.Conv2d(in_channels, in_channels, kernel_size = 1)
        self.output_conv = nn.Conv2d(in_channels, in_channels, kernel_size = 1)
    
    def forward(self, F):
        # first branch
        F1 = self.first_branch_conv(F)                  # (C/8, W, H)
        F1 = F1.reshape((F1.size(0), F1.size(1), -1))   # (C/8, W*H)
        F1 = torch.transpose(F1, -2, -1)                # (W*H, C/8)
        # second branch
        F2 = self.second_branch_conv(F)                 # (C/8, W, H)
        F2 = F2.reshape((F2.size(0), F2.size(1), -1))   # (C/8, W*H)
        F2 = nn.Softmax(dim = -1)(torch.matmul(F1, F2)) # (W*H, W*H)
        # third branch
        F3 = self.third_branch_conv(F)                  # (C, W, H)
        F3 = F3.reshape((F3.size(0), F3.size(1), -1))   # (C, W*H)
        F3 = torch.matmul(F3, F2)                       # (C, W*H)
        F3 = F3.reshape(F.shape)                        # (C, W, H)
        return self.output_conv(F3*F)

How to re-write the above code to include one more dimension "batch size"?
Thank you,
Ling

in Pytorch you dont have to specify anything in the model class for the batchsize. You just have to be sure that the input (in your case 'F') has the dimension (B, C, W, H), B is the batch-size. YOu can do that either by creating a tensor of images yourself, or using a dataloader. You have to create an own class for your dataset with this structure:

class YourDataset(torch.utils.data.Dataset):
    def __init__(self, datase):
        """ constructor
        load your dataset here,
        dataset could be a csv file or a list or whatever you want
        """

    def __getitem__(self, i)
        """ get sample from dataset
        :param int i: index of dataset (iterator of training-loop)
        :return tensor: preprocessed sample and target

        preprocess the i'th sample/image here if necessary and/or just return the image 
        and its label
        """
    def __len__(self):
        """ returns length of dataset """

Here you still dont specify anything about the batchsize, but that comes with the dataloader, which you can create like this:
dataset = YourDataset(dataset_file_or_whatever)
dataloder = torch.utils.data.DataLoader(
        dataset ,
        batch_size=64,
        num_workers=1,
        shuffle=True)

And know you can just iterate through the dataloder in your train loop, for example like this:
for epoch in range(epochs):
    for images, targets in dataloader:
        # images and targets will now have the dimensions (B, C, W, H)!
        # thats it! Now you can train with batches

I hope this answered your question!
Edit
Just saw that it didnt really, but I wrote all that stuff so imma let it stand here, hope its still helpful for those who dont know

Thank you @Theodor Peifer, I have got an idea what to do.
