Hi Prithviraj!



 Xanthan:

I’m performing global classification.


In that case, you will probably want to use CrossEntropyLoss as
your loss criterion.  Therefore you will not want to use an explicit
Softmax as CrossEntropyLoss has, in effect, Softmax built in.





 KFrank:

are willing to flatten the input tensor


Will there be a need to do that?


If you don’t want to have a non-linearity (the activation in my example,
below) then the convolution becomes redundant.  So you might as well
just have a single Linear layer (no Conv2d) for which you will need
to flatten the (non-batch dimensions of the) tensor you pass into the
Linear.  This is illustrated below.


by the time it comes to last layer, I believe the shape is [batch_size,64*8,4,4]
and I’m (want) converting this to shape  [batch_size,output_channels,1,1] after the convolution (then the linear and softmax stuff can come).
So now you know the shapes. what do u recommend i should do?


It really depends on what you want to do.  What is your proposed
Conv2d supposed to accomplish?  What is the final Linear supposed
to accomplish?
In any event, you can shrink your last two dimensions, [4, 4] to
[1, 1] by convolving with a kernel of size 4.
Then to pass this to a subsequent Linear you have to get rid of the
trailing [1, 1] dimensions (for example, by squeeze()ing).
These various points are illustrated by the following script:
import torch
print (torch.__version__)

_ = torch.manual_seed (2021)

n648 = 3   # smaller placeholder for 64*8
batch_size = 2
output_channels = 5
conv = torch.nn.Conv2d (n648, output_channels, 4)
activation = torch.nn.Sigmoid()   # for example
lin = torch.nn.Linear (output_channels, output_channels)

t1 = torch.randn (batch_size, n648, 4, 4)
print ('t1.shape =', t1.shape)
t2 = conv (t1)
print ('t2.shape =', t2.shape)
t3 = activation (t2)   # without non-linearity, conv becomes redundant
print ('t3.shape =', t3.shape)
t4 = t3.squeeze()
print ('t4.shape =', t4.shape)

prediction = lin (t4)
print ('prediction.shape =', prediction.shape)
target = torch.randint (output_channels, (batch_size,))
print ('target.shape =', target.shape)

# no softmax between final Linear layer and CrossEntropyLoss

loss_fn = torch.nn.CrossEntropyLoss()
loss = loss_fn (prediction, target)
print ('loss =', loss)

# if you want just a single Linear layer, you can flatten t1

t1_flat = torch.flatten (t1, start_dim = 1)
print ('t1_flat.shape =', t1_flat.shape)
lin2 = torch.nn.Linear (n648 * 4 * 4, output_channels)
prediction2 = lin2 (t1_flat)
print ('prediction2.shape =', prediction2.shape)
loss2 = loss_fn (prediction2, target)
print ('loss2 =', loss2)

And here is its output:
1.7.1
t1.shape = torch.Size([2, 3, 4, 4])
t2.shape = torch.Size([2, 5, 1, 1])
t3.shape = torch.Size([2, 5, 1, 1])
t4.shape = torch.Size([2, 5])
prediction.shape = torch.Size([2, 5])
target.shape = torch.Size([2])
loss = tensor(1.6289, grad_fn=<NllLossBackward>)
t1_flat.shape = torch.Size([2, 48])
prediction2.shape = torch.Size([2, 5])
loss2 = tensor(2.1311, grad_fn=<NllLossBackward>)

Best.
K. Frank