파이토치에서 드롭아웃 구현: 예제 포함

self.dropout = nn.Dropout(0.25)

class Net(nn.Module):
  def __init__(self, input_shape=(3,32,32)):
    super(Net, self).__init__()
    
    self.conv1 = nn.Conv2d(3, 32, 3)
    self.conv2 = nn.Conv2d(32, 64, 3)
    self.conv3 = nn.Conv2d(64, 128, 3)
    
    self.pool = nn.MaxPool2d(2,2)

    n_size = self._get_conv_output(input_shape)
    
    self.fc1 = nn.Linear(n_size, 512)
    self.fc2 = nn.Linear(512, 10)

    # Define proportion or neurons to dropout
    self.dropout = nn.Dropout(0.25)
      
  def forward(self, x):
    x = self._forward_features(x)
    x = x.view(x.size(0), -1)
    x = self.dropout(x)
    x = F.relu(self.fc1(x))
    # Apply dropout
    x = self.dropout(x)
    x = self.fc2(x)
    return x

def train(model, device, train_loader, optimizer, criterion, epoch, steps_per_epoch=20):

  # Log gradients and model parameters
  wandb.watch(model)

  # loop over the data iterator, and feed the inputs to the network and adjust the weights.
  for batch_idx, (data, target) in enumerate(train_loader, start=0):
    
    # ...
    
    acc = round((train_correct / train_total) * 100, 2)
    # Log metrics to visualize performance
    wandb.log({'Train Loss': train_loss/train_total, 'Train Accuracy': acc})