adjusted mrnet

mrnet/train.py

@@ -15,7 +15,7 @@ from utils.eval import eval_net
 
 dir_img = 'data/train_imgs/'
 dir_mask = 'data/train_masks/'
-dir_checkpoint = 'checkpoint'
+dir_checkpoint = 'checkpoint/'
 
 
 def train_net(net, device, epochs = 5, batch_size = 1, lr = 0.1, val_percent = 0.1, img_scale = 0.5):
@@ -32,6 +32,7 @@ def train_net(net, device, epochs = 5, batch_size = 1, lr = 0.1, val_percent = 0
     for epoch in tqdm(range(epochs)):
         net.train()
         epoch_loss = 0
+        with tqdm(total = n_train, desc = f'Epoch {epoch + 1}/{epochs}', unit = 'img') as pbar:
             for batch in train_loader:
                 imgs = batch['image']
                 true_masks = batch['mask']
@@ -43,10 +44,12 @@ def train_net(net, device, epochs = 5, batch_size = 1, lr = 0.1, val_percent = 0
                 masks_pred = net(imgs)
                 loss = criterion(masks_pred, true_masks)
                 epoch_loss += loss.item()
+                pbar.set_postfix(**{'loss (batch)': loss.item()})
 
                 optimizer.zero_grad()
                 loss.backward()
                 optimizer.step()
+                pbar.update(imgs.shape[0])
 
         val_score = eval_net(net, val_loader, device, n_val)
         logging.info('Validation cross entropy: {}'.format(val_score))