Skip to content

Instantly share code, notes, and snippets.

@level14taken
Created December 29, 2020 10:05
Show Gist options
  • Select an option

  • Save level14taken/2bd3514d0a0bc0bc8575bf2e0dd5ed51 to your computer and use it in GitHub Desktop.

Select an option

Save level14taken/2bd3514d0a0bc0bc8575bf2e0dd5ed51 to your computer and use it in GitHub Desktop.
def train_loop_fn(loader):
model.train()
loss_,iou_,acc_=0.,0.,0.
for x ,(image, mask) in enumerate(loader):
mask=mask.to(device)
image= image.to(device)
y_pred = model(image)
loss = loss_fn(y_pred, mask)
loss_+= loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
ACC=acc(y_pred,mask)
IOU= iou_metric(y_pred,mask)
acc_+=ACC.item()
iou_+=IOU.item()
if scheduler:
scheduler.step()
loss_/=len(loader)
iou_/=len(loader)
acc_/=(len(train)*101*101)
print('Train[{}]: Loss={:.5f} IOU={:.3f} ACC={:.3f}'.format(
x, loss_,iou_,acc_))
return loss_,iou_,acc_
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment