human
I am working on the file "Ass2.5/END3_MNIST_addnumber_ass2_5.ipynb".
The first task is:
Model
The model consists of 3 convolutional blocks, followed by 3 linear blocks, with ReLU activation in between. The MNIST image is fed to the 1st Conv block, and the ranbdom number (one hot) to the first linear block (alongwith the output of the 3rd Conv block, concatenated).
Can you write Python code for it?
 Write me a the Python code you requested?
assistant
#Here is the Python code you requested.

class Model(nn.Module):
  def __init__(self):
    super(Model, self).__init__()
    # conv layer 1
    self.conv1 = nn.Sequential(
                          
        nn.Conv2d(1,16,5), # 16x24x24
        nn.ReLU(),
        nn.MaxPool2d(2,2)  # 16x12x12
    )
    self.conv2 = nn.Sequential(
        nn.Conv2d(16,32,5), # 32x8x8
        nn.ReLU(),
        nn.MaxPool2d(2,2)  # 32x4x4
    )
    self.conv3 = nn.Sequential(
        nn.Conv2d(32,10,3), # 10x2x2
        nn.MaxPool2d(2,2)   # 10x1x1
    )
    self.relu = nn.ReLU()
    self.fc1 = nn.Linear(10+10, 60) # adding random number one hot to the 1x10 MNIST output
    self.fc2 = nn.Linear(60, 30)
    self.fc3 = nn.Linear(30, 15) # 10 for MNIST 1-hot coding, and 5 for binary repr of sum of digits

  def forward(self, image, number):
    #print("0 ",image.shape)
    x = self.conv1(image)
    #print("1 ",x.shape)
    x = self.conv2(x)
    #print("2 ",x.shape)
    x = self.conv3(x)
    #print("3 ",x.shape)
    x = x.view(-1,10)
    #print("after ",x.shape)
    # concatenate the number
    x = torch.cat((x,number),1)
    x = self.fc1(x)
    x = self.relu(x)
    x = self.fc2(x)
    x = self.relu(x)
    x = self.fc3(x)
    x = x.view(-1,15)
    #print("In forward x shape ",x.shape)
    
    # The first 10 outputs should be the onehot encoding of the MNIST digit
    # using a Log softmax (with NLL Loss) for this
    o1 = F.log_softmax(x[:,:10])
    #print("In forward o1 shape ",o1.shape)

    # for the 5 digit sum outout - as it is a multi-label classification, I am using a Sigmoid and not a softmax as there
    # will be multiple 1's in the output
    # used Hardsigmoid as it has a more sharp curve
    sig = nn.Hardsigmoid()
    o2 = sig(x[:,10:])
    #print("In forward o2 shape ",o2.shape)
    return torch.cat((o1,o2),1)

model = Model().to(device)
print(model)
# random test
model.forward(torch.rand((1, 1, 28, 28)).to(device), torch.rand((1, 10)).to(device))

def train(model, device, train_loader, optimizer, epoch, losses):
    print(f"EPOCH - {epoch}")
   ...