Train model

[zhouzhenghong-gt]

Train model

[Sungden]

Hi! this is a great paper! and I want to ask how to use trainer.py to train the model?

Hello, Do you know how to train the model yet? Thanks.

[sainatarajan]

@Sungden If you're still working on this, the below code should help you get a start

import torch
import torch.nn as nn
import argparse

from torchsummary import summary
from net import *
from trainer import Trainer_ReconNet
from data import MyReconDataset # Your dataset class

import os
# os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'

###########################
########ARGUMENTS##########

parser = argparse.ArgumentParser(description='A test program.')

parser.add_argument("--exp", default='exp1', help="test1")
parser.add_argument("--arch", default='ReconNet', help="test1")
parser.add_argument("--print_freq", default=2, help="test1")
parser.add_argument("--output_path", default='output/', help="test1")
parser.add_argument("--resume", default=False, help="test1")
parser.add_argument("--loss", default='l2', help="test1")
parser.add_argument("--optim", default='adam', help="test1")
parser.add_argument("--lr", default=1e-4, help="test1")
parser.add_argument("--weight_decay", default=0.99, help="test1")

args = parser.parse_args()

###########################

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

trainer_obj = Trainer_ReconNet(args)

root_path = '/path/to/your/dataset'
train_dataset = MyReconDataset(root_path + '/images/*.png', root_path + '/volumes/*.nii.gz')
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=1)

for epoch in range(5):
    print('Epoch:', epoch+1)

    trainer_obj.train_epoch(train_dataloader, epoch + 1)

[yinyin-llll]

hi，i also want to retrain the model but it need the csv file with annotation.Do you know this file or how to design this file?
Thank you very much!