model = LeNet()
if loadState:
model.load_state_dict(torch.load("bestState.pth"))
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
loss_fn = nn.CrossEntropyLoss()
if Training:
while model.EpochRunner:
trainloader = torch.utils.data.DataLoader(TrainingData, batch_size=250, shuffle=True, num_workers=8)
for batch in trainloader:
optimizer.zero_grad()
batchLabel = batch[:,0]
batchData = batch[:,1:].reshape(250,1,28,28).float()
y_pred = model.forward(batchData)
loss = loss_fn(y_pred, batchLabel)
loss.backward()
optimizer.step()
model.CurrentValidationLoss.clear()
validationLoader = torch.utils.data.DataLoader(ValidationData, batch_size=250, shuffle=True, num_workers=8)
for batch in validationLoader:
batchLabel = batch[:,0]
batchData = batch[:,1:].reshape(250,1,28,28).float()
y_pred = model.forward(batchData)
loss = loss_fn(y_pred, batchLabel)
model.CurrentValidationLoss.append(float(loss.item()))
model.OverallValidationLoss.append(np.mean(np.array(model.CurrentValidationLoss)))
