def Train(train_data, val_data, net, num_epoch, lr, output_dir):
net = torch.nn.DataParallel(net, device_ids=[0])
solver = Solver(net, train_data, 0.0001, output_dir)
solver.criterion = lambda p, t: SegLoss(
p, t, num_classes=2, loss_fn=DiceLoss)
solver.iter_per_sample = 100
for i_epoch in range(0, num_epoch, solver.iter_per_sample):
# train
solver.dataset.set_trans_prob(i_epoch / 1000.0 + 0.15)
loss = solver.step_one_epoch(batch_size=2, iter_size=1)
i_epoch = solver.num_epoch
print(('epoch:%d, loss:%f') % (i_epoch, loss))
if i_epoch % 100 == 0:
save_path = solver.save_model()
print('save model at %s' % save_path)
# val
if i_epoch % 100 == 0:
'''
print('val')
eval_dict_val = Evaluate(net, val_data, 'val')
for key, value in eval_dict_val.items():
solver.writer.add_scalar(key, value, i_epoch)
'''
'''
def Train(train_data, val_data, net, num_epoch, lr, output_dir):
net = torch.nn.DataParallel(net, device_ids=[0, 1, 2, 3])
solver = Solver(net, train_data, 0.0001, output_dir)
solver.criterion = lambda p, t: SegLoss(
p, t, num_classes=5, loss_fn=FocalLoss(5))
solver.iter_per_sample = 100
for i_epoch in range(0, num_epoch, solver.iter_per_sample):
# train
solver.dataset.set_trans_prob(1.0)
loss = solver.step_one_epoch(batch_size=36, iter_size=1)
i_epoch = solver.num_epoch
print(('epoch:%d, loss:%f') % (i_epoch, loss))
if i_epoch % 100 == 0:
save_path = solver.save_model()
print('save model at %s' % save_path)
# val
'''
def Train(train_set, val_set, net, num_epoch, lr, output_dir):
solver = Solver(net, train_set, 0.0001, output_dir)
solver.criterion = lambda p, t: SegLoss(
p, t, num_classes=5, loss_fn=FocalLoss(
5)) # pred, target, num_classes set here
solver.iter_per_sample = 100 # goes to self.dataset.set_iter_per_sample() in Solver
for i_epoch in range(0, num_epoch, solver.iter_per_sample
): # say, (0, 2000, 100), so i_epoch = 0, 100, 200...
# train
solver.dataset.set_trans_prob(i_epoch / 1000.0 + 0.15)
loss = solver.step_one_epoch(batch_size=10, iter_size=1)
i_epoch = solver.num_epoch # num_epoch is accumulated in step_one_epoch by iter_per_sample
print('Epoch: %d, Loss: %f' % (i_epoch, loss))
if i_epoch % 100 == 0:
save_path = solver.save_model()
print('Save model at %s' % save_path)
# val
if i_epoch % 100 == 0:
eval_dict_val = Evaluate(net, val_set, 'val')
for key, value in eval_dict_val.items():
solver.writer.add_scalar(key, value, i_epoch)