model.to(device)
model = torch.nn.DataParallel(model, device_ids=FG.devices)
#fname = 'Plane2D_79'
fname = 'plane-5crop'
fname = 'Plane2D_79_' + str(FG.axis)
if FG.isize == 157:
fname = 'Plane2D_157_' + str(FG.axis)
load_checkpoint(FG.checkpoint_root,
FG.running_fold,
fname,
model.module,
epoch=None,
is_best=True)
criterion = torch.nn.CrossEntropyLoss().to(device)
G = sliceGenerator(FG).to('cuda:{}'.format(FG.devices[0]))
G = torch.nn.DataParallel(G, FG.devices)
if FG.isize == 79:
G.load_state_dict(
torch.load(os.path.join('result_2d_classify', 'G_79.pth')))
else:
G.load_state_dict(
torch.load(os.path.join('result_2d_classify', 'G.pth')))
optimizerG = optim.Adam(G.parameters(), lr=FG.lrD, betas=(0.5, 0.999))
# make dataset
x, y, train_idx, test_idx, ratio = fold_split(FG)
# original dataset
testset = ADNIDataset2D(
FG,
x[test_idx],
transform=Compose([ToWoldCoordinateSystem(), ToTensor()])
else :
transform=Compose([ToWoldCoordinateSystem(), Normalize(0.2,0.9), ToTensor()])
trainset = ADNIDataset(FG, x, y, transform=transform)
trainloader = DataLoader(trainset, batch_size=FG.batch_size,
shuffle=True, pin_memory=True,
num_workers=4)
for _, data in enumerate(trainloader):
data = data['image']
break
if '3d' in FG.model:
G = Generator(FG).to('cuda:{}'.format(FG.devices[0]))
D = Discriminator(FG).to('cuda:{}'.format(FG.devices[0]))
elif '2d' in FG.model:
G = sliceGenerator(FG, output_dim=data.shape[1]).to('cuda:{}'.format(FG.devices[0]))
D = sliceDiscriminator(FG, input_dim=data.shape[1]).to('cuda:{}'.format(FG.devices[0]))
elif 'resnet' in FG.model:
G = ResGenerator(FG).to('cuda:{}'.format(FG.devices[0]))
D = ResDiscriminator(FG).to('cuda:{}'.format(FG.devices[0]))
elif 'mix' in FG.model:
G = ResGenerator(FG).to('cuda:{}'.format(FG.devices[0]))
D = Discriminator(FG).to('cuda:{}'.format(FG.devices[0]))
elif 'bigi' in FG.model:
G = bigiGenerator(FG).to('cuda:{}'.format(FG.devices[0]))
D = bigiDiscriminator(FG).to('cuda:{}'.format(FG.devices[0]))
elif '2dmix-1' in FG.model:
G = ResGenerator2D(FG, output_dim=data.shape[1]).to('cuda:{}'.format(FG.devices[0]))
D = bigiDiscriminator(FG, input_dim=data.shape[1]).to('cuda:{}'.format(FG.devices[0]))
elif '2dmix-2' in FG.model:
G = ResGenerator2D(FG, output_dim=data.shape[1]).to('cuda:{}'.format(FG.devices[0]))