現(xiàn)在，我們將定義我們的超參數(shù)

# Hyper Parameters

input_size = 784

num_classes = 10

num_epochs = 5

batch_size = 100

learning_rate = 0.001

在我們的數(shù)據(jù)集中，圖像大小為28 * 28。因此，我們的輸入大小是784.此外，這里有10位數(shù)字，因此，我們可以有10個不同的輸出。因此，我們將num_classes設置為10.此外，我們將在整個數(shù)據(jù)集上訓練五次。最后，我們將分別訓練小批量的100張圖像，以防止因內存溢出而導致程序崩潰。

在此之后，我們將定義我們的模型如下。在這里，我們將我們的模型初始化為torch.nn.Module的子類，然后定義前向傳遞。在我們編寫的代碼中，softmax在每次正向傳遞期間內部計算，因此我們不需要在forward（）函數(shù)內指定它。

class LogisticRegression(nn.Module):

def __init__(self, input_size, num_classes):

super(LogisticRegression, self).__init__()

self.linear = nn.Linear(input_size, num_classes)

def forward(self, x):

out = self.linear(x)

return out

定義了我們的類之后，現(xiàn)在我們實例化了一個對象

model = LogisticRegression(input_size, num_classes)

接下來，我們設置損失函數(shù)和優(yōu)化器。在這里，我們將使用交叉熵損失，對于優(yōu)化器，我們將使用隨機梯度下降算法，其學習率為0.001，如上面的超參數(shù)中所定義。

criterion = nn.CrossEntropyLoss()

optimizer = torch.optim.SGD(model.parameters(), lr = learning_rate)