Beispiel #1
0
def test():
    model = ModelFlow_stride(2, 3, opt.start_channel).cuda()
    transform = SpatialTransform().cuda()

    model.load_state_dict(torch.load(opt.modelpath))
    model.eval()
    transform.eval()

    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(grid, (1, ) +
                                                grid.shape))).cuda().float()

    start = timeit.default_timer()

    A = Variable(torch.from_numpy(load_5D(opt.fixed))).cuda().float()
    B = Variable(torch.from_numpy(load_5D(opt.moving))).cuda().float()
    start2 = timeit.default_timer()
    print('Time for loading data: ', start2 - start)
    pred = model(A, B)
    F_AB = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    F_AB = F_AB.astype(np.float32) * range_flow
    warped_A = transform(A,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    start3 = timeit.default_timer()
    print('Time for registration: ', start3 - start2)

    warped_F_BA = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    warped_F_BA = warped_F_BA.astype(np.float32) * range_flow

    start4 = timeit.default_timer()
    print('Time for generating inverse flow: ', start4 - start3)

    save_flow(F_AB, savepath + '/flow_A_B.nii.gz')
    save_flow(warped_F_BA, savepath + '/inverse_flow_B_A.nii.gz')
    save_img(warped_A, savepath + '/warped_A.nii.gz')

    start5 = timeit.default_timer()
    print('Time for saving results: ', start5 - start4)
    del pred

    pred = model(B, A)
    F_BA = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    F_BA = F_BA.astype(np.float32) * range_flow
    warped_B = transform(B,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    warped_F_AB = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    warped_F_AB = warped_F_AB.astype(np.float32) * range_flow
    save_flow(F_BA, savepath + '/flow_B_A.nii.gz')
    save_flow(warped_F_AB, savepath + '/inverse_flow_A_B.nii.gz')
    save_img(warped_B, savepath + '/warped_B.nii.gz')
Beispiel #2
0
def train():
    model =ModelFlow_stride(2,3,start_channel).cuda()
    loss_similarity =mse_loss
    loss_inverse = mse_loss
    loss_antifold = antifoldloss
    loss_smooth = smoothloss
    transform = SpatialTransform().cuda()
    for param in transform.parameters():
        param.requires_grad = False
        param.volatile=True
    names = glob.glob(datapath + '/*.gz')
    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(grid, (1,) + grid.shape))).cuda().float()

    print(grid.type())
    optimizer = torch.optim.Adam(model.parameters(),lr=lr) 
    model_dir = '../Model'

    if not os.path.isdir(model_dir):
        os.mkdir(model_dir)

    lossall = np.zeros((5,iteration))
    training_generator = Data.DataLoader(Dataset(names,iteration,True), batch_size=1,
                        shuffle=False, num_workers=2)
    step=0
    for  X,Y in training_generator:

        X = X.cuda().float()
        Y = Y.cuda().float()
        F_xy = model(X,Y)
        F_yx = model(Y,X)
    
        X_Y = transform(X,F_xy.permute(0,2,3,4,1)*range_flow,grid)
        Y_X = transform(Y,F_yx.permute(0,2,3,4,1)*range_flow,grid)
        # Note that, the generation of inverse flow depends on the definition of transform. 
        # The generation strategies are sligtly different for the backward warpping and forward warpping
        F_xy_ = transform(-F_yx,F_xy.permute(0,2,3,4,1)*range_flow,grid)
        F_yx_ = transform(-F_xy,F_yx.permute(0,2,3,4,1)*range_flow,grid)
        loss1 = loss_similarity(Y,X_Y) + loss_similarity(X,Y_X)
        loss2 = loss_inverse(F_xy*range_flow,F_xy_*range_flow) + loss_inverse(F_yx*range_flow,F_yx_*range_flow)
        
        
        loss3 =  loss_antifold(F_xy*range_flow) + loss_antifold(F_yx*range_flow)
        loss4 =  loss_smooth(F_xy*range_flow) + loss_smooth(F_yx*range_flow)
        loss = loss1+inverse*loss2 + antifold*loss3 + smooth*loss4
        optimizer.zero_grad()           # clear gradients for this training step
        loss.backward()                 # backpropagation, compute gradients
        optimizer.step()                # apply gradients
        lossall[:,step] = np.array([loss.item(),loss1.item(),loss2.item(),loss3.item(),loss4.item()])
        sys.stdout.write("\r" + 'step "{0}" -> training loss "{1:.4f}" - sim "{2:.4f}" - inv "{3:.4f}" \
            - ant "{4:.4f}" -smo "{5:.4f}" '.format(step, loss.item(),loss1.item(),loss2.item(),loss3.item(),loss4.item()))
        sys.stdout.flush()
        if(step % n_checkpoint == 0):
            modelname = model_dir + '/' + str(step) + '.pth'
            torch.save(model.state_dict(), modelname)
        step+=1
    np.save(model_dir+'/loss.npy',lossall)
def test():
    model = SYMNet(2, 3, opt.start_channel).cuda()
    transform = SpatialTransform().cuda()

    diff_transform = DiffeomorphicTransform(time_step=7).cuda()
    com_transform = CompositionTransform().cuda()

    model.load_state_dict(torch.load(opt.modelpath))
    model.eval()
    transform.eval()
    diff_transform.eval()
    com_transform.eval()

    grid = generate_grid(imgshape)
    grid = torch.from_numpy(np.reshape(grid,
                                       (1, ) + grid.shape)).cuda().float()

    use_cuda = True
    device = torch.device("cuda" if use_cuda else "cpu")

    fixed_img = load_4D(fixed_path)
    moved_img = load_4D(moving_path)

    fixed_img = torch.from_numpy(fixed_img).float().to(device).unsqueeze(dim=0)
    moved_img = torch.from_numpy(moved_img).float().to(device).unsqueeze(dim=0)

    with torch.no_grad():
        F_xy, F_yx = model(fixed_img, moved_img)

        F_X_Y_half = diff_transform(F_xy, grid, range_flow)
        F_Y_X_half = diff_transform(F_yx, grid, range_flow)

        F_X_Y_half_inv = diff_transform(-F_xy, grid, range_flow)
        F_Y_X_half_inv = diff_transform(-F_yx, grid, range_flow)

        F_X_Y = com_transform(F_X_Y_half, F_Y_X_half_inv, grid, range_flow)
        F_Y_X = com_transform(F_Y_X_half, F_X_Y_half_inv, grid, range_flow)

        F_BA = F_Y_X.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
        F_BA = F_BA.astype(np.float32) * range_flow

        F_AB = F_X_Y.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
        F_AB = F_AB.astype(np.float32) * range_flow

        warped_B = transform(moved_img,
                             F_Y_X.permute(0, 2, 3, 4, 1) * range_flow,
                             grid).data.cpu().numpy()[0, 0, :, :, :]
        warped_A = transform(fixed_img,
                             F_X_Y.permute(0, 2, 3, 4, 1) * range_flow,
                             grid).data.cpu().numpy()[0, 0, :, :, :]

        save_flow(F_BA, savepath + '/wrapped_flow_B_to_A.nii.gz')
        save_flow(F_AB, savepath + '/wrapped_flow_A_to_B.nii.gz')
        save_img(warped_B, savepath + '/wrapped_norm_B_to_A.nii.gz')
        save_img(warped_A, savepath + '/wrapped_norm_A_to_B.nii.gz')

        print("Finished.")
def train():
    torch.cuda.empty_cache()
    device = torch.device("cuda:6")
    model = ModelFlow_stride(2, 3, start_channel).cuda(device)
    # model =ModelFlow_stride(2,3,start_channel).cpu()
    loss_similarity = mse_loss
    loss_inverse = mse_loss
    loss_antifold = antifoldloss
    loss_smooth = smoothloss
    transform = SpatialTransform().cuda(device)
    # transform = SpatialTransform().cpu()
    for param in transform.parameters():
        param.requires_grad = False
        param.volatile = True
    names = glob.glob(datapath + '/*.gz')
    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(
        grid, (1, ) + grid.shape))).cuda(device).float()
    # grid = Variable(torch.from_numpy(np.reshape(grid, (1,) + grid.shape))).cpu().float()

    print(grid.type())
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    model_dir = '../Model'

    if not os.path.isdir(model_dir):
        os.mkdir(model_dir)

    lossall = np.zeros((5, iteration))
    training_generator = Data.DataLoader(Dataset(names, iteration, True),
                                         batch_size=1,
                                         shuffle=False,
                                         num_workers=0)
    step = 0
    for X, Y in training_generator:
        X = X.cuda(device).float()
        # X = X.cpu().float()
        Y = Y.cuda(device).float()
        # Y = Y.cpu().float()
        # # X = sitk.GetArrayFromImage(sitk.ReadImage(X, sitk.sitkFloat32))
        # # Y = sitk.GetArrayFromImage(sitk.ReadImage(Y, sitk.sitkFloat32))
        #X = Variable(X).cuda(device).float()
        #Y = Variable(Y).cuda(device).float()
        X, Y = train_padding(X, Y)  #added today

        X = X.cuda(device).float(
        )  #added as the values returned after padding was not cuda
        Y = Y.cuda(device).float()

        F_xy = model(X, Y)
        F_yx = model(Y, X)
        X_Y = transform(X, F_xy.permute(0, 2, 3, 4, 1) * range_flow, grid)
        Y_X = transform(Y, F_yx.permute(0, 2, 3, 4, 1) * range_flow, grid)
        F_xy_ = transform(-F_xy,
                          F_xy.permute(0, 2, 3, 4, 1) * range_flow, grid)
        F_yx_ = transform(-F_yx,
                          F_yx.permute(0, 2, 3, 4, 1) * range_flow, grid)
        loss1 = loss_similarity(Y, X_Y) + loss_similarity(X, Y_X)
        loss2 = loss_inverse(F_xy * range_flow,
                             F_xy_ * range_flow) + loss_inverse(
                                 F_yx * range_flow, F_yx_ * range_flow)
        loss3 = loss_antifold(F_xy * range_flow) + loss_antifold(
            F_yx * range_flow)
        loss4 = loss_smooth(F_xy * range_flow) + loss_smooth(F_yx * range_flow)
        loss = loss1 + inverse * loss2 + antifold * loss3 + smooth * loss4
        optimizer.zero_grad()  # clear gradients for this training step
        loss.backward()  # backpropagation, compute gradients
        optimizer.step()  # apply gradients
        #lossall[:,step] = np.array([loss.data[0],loss1.data[0],loss2.data[0],loss3.data[0],loss4.data[0]])
        lossall[:, step] = np.array(
            [loss.data, loss1.data, loss2.data, loss3.data, loss4.data])
        #sys.stdout.write("\r" + 'step "{0}" -> training loss "{1:.4f}" - sim "{2:.4f}" - inv "{3:.4f}" - ant "{4:.4f}" -smo "{5:.4f}" '.format(step, loss.data[0],loss1.data[0],loss2.data[0],loss3.data[0],loss4.data[0]))
        sys.stdout.write(
            "\r" +
            'step "{0}" -> training loss "{1:.4f}" - sim "{2:.4f}" - inv "{3:.4f}" - ant "{4:.4f}" -smo "{5:.4f}" '
            .format(step, loss.data, loss1.data, loss2.data, loss3.data,
                    loss4.data))
        sys.stdout.flush()
        if (step % n_checkpoint == 0):
            modelname = model_dir + '/' + str(step) + '.pth'
            torch.save(model.state_dict(), modelname)
        step += 1
        torch.cuda.empty_cache()
        del X
        del Y
    np.save(model_dir + '/loss.npy', lossall)
Beispiel #5
0
def test():
    device = torch.device("cuda:6")
    model = ModelFlow_stride(2, 3, opt.start_channel).cuda(device)
    # model =ModelFlow_stride(2,3,opt.start_channel).cpu()
    # transform = SpatialTransform().cpu()

    transform = SpatialTransform().cuda(device)
    # model.load_state_dict(torch.load(opt.modelpath))
    # model.load_state_dict(torch.load(opt.modelpath, map_location=torch.device('cpu')))
    model.load_state_dict(torch.load(opt.modelpath, map_location=device))

    model.eval()
    transform.eval()
    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(
        grid, (1, ) + grid.shape))).cuda(device).float()
    # grid = Variable(torch.from_numpy(np.reshape(grid, (1,) + grid.shape))).cpu().float()

    start = timeit.default_timer()
    A = sitk.GetArrayFromImage(sitk.ReadImage(opt.fixed, sitk.sitkFloat32))
    B = sitk.GetArrayFromImage(sitk.ReadImage(opt.moving, sitk.sitkFloat32))

    #A = Variable(torch.from_numpy(A)).cuda(device).float()
    #B = Variable(torch.from_numpy(B)).cuda(device).float()

    # A, B = padding(A,B)
    #A = load_5D(A).cuda(device).float()
    #B = load_5D(B).cuda(device).float()

    A = load_5D(A)
    B = load_5D(B)

    A = Variable(torch.from_numpy(A)).cuda(device).float()
    B = Variable(torch.from_numpy(B)).cuda(device).float()

    #A = Variable(torch.from_numpy( load_5D(opt.fixed))).cuda(device).float()
    #B = Variable(torch.from_numpy( load_5D(opt.moving))).cuda(device).float()
    start2 = timeit.default_timer()
    print('Time for loading data: ', start2 - start)
    pred = model(A, B)
    F_AB = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    #F_AB = pred.permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]

    F_AB = F_AB.astype(np.float32) * range_flow
    warped_A = transform(A,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    #warped_A = transform(A,pred.permute(0,2,3,4,1)*range_flow,grid).data.cuda(device).numpy()[0, 0, :, :, :]
    start3 = timeit.default_timer()
    print('Time for registration: ', start3 - start2)
    warped_F_BA = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    #warped_F_BA = transform(-pred,pred.permute(0,2,3,4,1)*range_flow,grid).permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]
    warped_F_BA = warped_F_BA.astype(np.float32) * range_flow
    start4 = timeit.default_timer()
    print('Time for generating inverse flow: ', start4 - start3)
    save_flow(F_AB, savepath + '/flow_A_B.nii.gz')
    save_flow(warped_F_BA, savepath + '/inverse_flow_B_A.nii.gz')
    save_img(warped_A, savepath + '/warped_A.nii.gz')
    start5 = timeit.default_timer()
    print('Time for saving results: ', start5 - start4)
    del pred
    pred = model(B, A)
    F_BA = pred.permute(0, 2, 3, 4, 1).data.cpu().numpy()[0, :, :, :, :]
    #F_BA = pred.permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]
    F_BA = F_BA.astype(np.float32) * range_flow
    warped_B = transform(B,
                         pred.permute(0, 2, 3, 4, 1) * range_flow,
                         grid).data.cpu().numpy()[0, 0, :, :, :]
    #warped_B = transform(B,pred.permute(0,2,3,4,1)*range_flow,grid).data.cuda(device).numpy()[0, 0, :, :, :]
    warped_F_AB = transform(-pred,
                            pred.permute(0, 2, 3, 4, 1) * range_flow,
                            grid).permute(0, 2, 3, 4,
                                          1).data.cpu().numpy()[0, :, :, :, :]
    #warped_F_AB = transform(-pred,pred.permute(0,2,3,4,1)*range_flow,grid).permute(0,2,3,4,1).data.cuda(device).numpy()[0, :, :, :, :]
    warped_F_AB = warped_F_AB.astype(np.float32) * range_flow
    save_flow(F_BA, savepath + '/flow_B_A.nii.gz')
    save_flow(warped_F_AB, savepath + '/inverse_flow_A_B.nii.gz')
    save_img(warped_B, savepath + '/warped_B.nii.gz')
Beispiel #6
0
def train():
    model = ModelFlow_stride(2, 3, start_channel).cuda()
    loss_similarity = NCC().loss
    loss_cycle = l1_loss
    loss_smooth = smoothloss
    transform = SpatialTransform().cuda()
    for param in transform.parameters():
        param.requires_grad = False
        param.volatile = True
    names = glob.glob(datapath + '/*.gz')
    grid = generate_grid(imgshape)
    grid = Variable(torch.from_numpy(np.reshape(grid, (1, ) +
                                                grid.shape))).cuda().float()

    print(grid.type())
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    model_dir = '../Model'

    if not os.path.isdir(model_dir):
        os.mkdir(model_dir)
    lossall = np.zeros((5, iteration))
    training_generator = Data.DataLoader(Dataset(names, iteration, False),
                                         batch_size=1,
                                         shuffle=False,
                                         num_workers=2)
    step = 0
    for X, Y in training_generator:

        X = X.cuda().float()
        Y = Y.cuda().float()
        F_xy = model(X, Y)
        F_yx = model(Y, X)

        X_Y = transform(X, F_xy.permute(0, 2, 3, 4, 1) * range_flow, grid)
        Y_X = transform(Y, F_yx.permute(0, 2, 3, 4, 1) * range_flow, grid)

        F_xy_ = model(Y_X, X_Y)
        F_yx_ = model(X_Y, Y_X)

        Y_X_Y = transform(Y_X, F_xy_.permute(0, 2, 3, 4, 1) * range_flow, grid)
        X_Y_X = transform(X_Y, F_yx_.permute(0, 2, 3, 4, 1) * range_flow, grid)

        F_xx = model(X, X)
        F_yy = model(Y, Y)
        X_X = transform(X, F_xx.permute(0, 2, 3, 4, 1) * range_flow, grid)
        Y_Y = transform(Y, F_yy.permute(0, 2, 3, 4, 1) * range_flow, grid)

        L_smooth = loss_smooth(F_xy * range_flow) + loss_smooth(
            F_yx * range_flow)
        L_regist = loss_similarity(Y, X_Y) + loss_similarity(X, Y_X) + \
                      lambda_ * L_smooth
        L_cycle = loss_cycle(X, X_Y_X) + loss_cycle(Y, Y_X_Y)
        L_identity = loss_similarity(X, X_X) + loss_similarity(Y, Y_Y)
        loss = L_regist + alpha * L_cycle + beta * L_identity

        optimizer.zero_grad()  # clear gradients for this training step
        loss.backward()  # backpropagation, compute gradients
        optimizer.step()  # apply gradients
        lossall[:, step] = np.array([
            loss.data[0], L_regist.data[0], L_cycle.data[0],
            L_identity.data[0], L_smooth.data[0]
        ])
        sys.stdout.write(
            "\r" +
            'step "{0}" -> training loss "{1:.4f}" - reg "{2:.4f}" - cyc "{3:.4f}" - ind "{4:.4f}" -smo "{5:.4f}" '
            .format(step, loss.data[0], L_regist.data[0], L_cycle.data[0],
                    L_identity.data[0], L_smooth.data[0]))
        sys.stdout.flush()
        if (step % n_checkpoint == 0):
            modelname = model_dir + '/' + str(step) + '.pth'
            torch.save(model.state_dict(), modelname)
        step += 1
    np.save(model_dir + '/loss.npy', lossall)
Beispiel #7
0
def train():
    model = SYMNet(2, 3, start_channel).cuda()
    loss_similarity = NCC()
    loss_smooth = smoothloss
    loss_magnitude = magnitude_loss
    loss_Jdet = neg_Jdet_loss

    transform = SpatialTransform().cuda()
    diff_transform = DiffeomorphicTransform(time_step=7).cuda()
    com_transform = CompositionTransform().cuda()

    for param in transform.parameters():
        param.requires_grad = False
        param.volatile = True
    names = sorted(glob.glob(datapath + '/*.nii'))[0:255]
    grid = generate_grid(imgshape)
    grid = torch.from_numpy(np.reshape(grid,
                                       (1, ) + grid.shape)).cuda().float()

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
    # optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9)
    model_dir = '../Model'

    if not os.path.isdir(model_dir):
        os.mkdir(model_dir)

    lossall = np.zeros((6, iteration))

    training_generator = Data.DataLoader(Dataset_epoch(names, norm=False),
                                         batch_size=1,
                                         shuffle=True,
                                         num_workers=2)
    step = 0

    while step <= iteration:
        for X, Y in training_generator:

            X = X.cuda().float()
            Y = Y.cuda().float()
            F_xy, F_yx = model(X, Y)

            F_X_Y_half = diff_transform(F_xy, grid, range_flow)
            F_Y_X_half = diff_transform(F_yx, grid, range_flow)

            F_X_Y_half_inv = diff_transform(-F_xy, grid, range_flow)
            F_Y_X_half_inv = diff_transform(-F_yx, grid, range_flow)

            X_Y_half = transform(
                X,
                F_X_Y_half.permute(0, 2, 3, 4, 1) * range_flow, grid)
            Y_X_half = transform(
                Y,
                F_Y_X_half.permute(0, 2, 3, 4, 1) * range_flow, grid)

            F_X_Y = com_transform(F_X_Y_half, F_Y_X_half_inv, grid, range_flow)
            F_Y_X = com_transform(F_Y_X_half, F_X_Y_half_inv, grid, range_flow)

            X_Y = transform(X, F_X_Y.permute(0, 2, 3, 4, 1) * range_flow, grid)
            Y_X = transform(Y, F_Y_X.permute(0, 2, 3, 4, 1) * range_flow, grid)

            loss1 = loss_similarity(X_Y_half, Y_X_half)
            loss2 = loss_similarity(Y, X_Y) + loss_similarity(X, Y_X)
            loss3 = loss_magnitude(F_X_Y_half * range_flow,
                                   F_Y_X_half * range_flow)
            loss4 = loss_Jdet(
                F_X_Y.permute(0, 2, 3, 4, 1) * range_flow, grid) + loss_Jdet(
                    F_Y_X.permute(0, 2, 3, 4, 1) * range_flow, grid)
            loss5 = loss_smooth(F_xy * range_flow) + loss_smooth(
                F_yx * range_flow)

            loss = loss1 + loss2 + magnitude * loss3 + local_ori * loss4 + smooth * loss5
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            lossall[:, step] = np.array([
                loss.item(),
                loss1.item(),
                loss2.item(),
                loss3.item(),
                loss4.item(),
                loss5.item()
            ])
            sys.stdout.write(
                "\r" +
                'step "{0}" -> training loss "{1:.4f}" - sim_mid "{2:.4f}" - sim_full "{3:4f}" - mag "{4:.4f}" - Jdet "{5:.10f}" -smo "{6:.4f}" '
                .format(step, loss.item(), loss1.item(), loss2.item(),
                        loss3.item(), loss4.item(), loss5.item()))
            sys.stdout.flush()

            if (step % n_checkpoint == 0):
                modelname = model_dir + '/SYMNet_' + str(step) + '.pth'
                torch.save(model.state_dict(), modelname)
                np.save(model_dir + '/loss_SYMNet_' + str(step) + '.npy',
                        lossall)
            step += 1

            if step > iteration:
                break
        print("one epoch pass")
    np.save(model_dir + '/loss_SYMNet.npy', lossall)