def main():
# import network architecture
builder = ModelBuilder()
model = builder.build_net(arch=TrainGlobalConfig.id,
num_input=TrainGlobalConfig.num_input,
num_classes=TrainGlobalConfig.num_classes,
num_branches=TrainGlobalConfig.num_branches,
padding_list=TrainGlobalConfig.padding_list,
dilation_list=TrainGlobalConfig.dilation_list)
model = model.to(device)
ch = torch.load(
f"./result/{TrainGlobalConfig.base_dir}/last-checkpoint.bin")
model.load_state_dict(fix_model_state_dict(ch["model_state_dict"]))
valid_list = os.path.join(TrainGlobalConfig.root_path, "valid_0.txt")
valid_dataset = BraTSDataset(
list_file=valid_list,
root=TrainGlobalConfig.root_path,
phase="val",
)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=1,
shuffle=False,
num_workers=TrainGlobalConfig.num_workers,
pin_memory=True)
test(valid_loader, model, TrainGlobalConfig)
def main():
# import network architecture
builder = ModelBuilder()
model = builder.build_net(arch=TrainGlobalConfig.id,
num_input=TrainGlobalConfig.num_input,
num_classes=TrainGlobalConfig.num_classes,
num_branches=TrainGlobalConfig.num_branches,
padding_list=TrainGlobalConfig.padding_list,
dilation_list=TrainGlobalConfig.dilation_list)
model = model.to(device)
optimizer = optim.RMSprop(model.parameters(),
TrainGlobalConfig.lr,
alpha=0.9,
eps=10**(-4),
weight_decay=1e-4,
momentum=0.6)
criterion = nn.CrossEntropyLoss()
train_list = os.path.join(TrainGlobalConfig.root_path, "train_0.txt")
valid_list = os.path.join(TrainGlobalConfig.root_path, "valid_0.txt")
train_dataset = BraTSDataset(list_file=train_list,
root=TrainGlobalConfig.root_path,
crop_size=TrainGlobalConfig.crop_size,
num_input=TrainGlobalConfig.num_input)
valid_dataset = BraTSDataset(list_file=valid_list,
root=TrainGlobalConfig.root_path,
crop_size=TrainGlobalConfig.crop_size,
num_input=TrainGlobalConfig.num_input)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=TrainGlobalConfig.batch_size,
shuffle=True,
num_workers=TrainGlobalConfig.num_workers,
pin_memory=True)
valid_loader = torch.utils.data.DataLoader(
valid_dataset,
batch_size=TrainGlobalConfig.batch_size,
shuffle=False,
num_workers=TrainGlobalConfig.num_workers,
pin_memory=True)
trainer = PytorchTrainer(
model=model,
optimizer=optimizer,
criterion=criterion,
device=device,
config=TrainGlobalConfig,
)
trainer.fit(train_loader, valid_loader)