def __init__(self, classes_num=3):
classes = classes_num
super(Net_3d_ALL, self).__init__()
self.in_block = NetInBlock(2, 24, 2)
self.down_block1 = NetDownBlock(24, 48, 2) #48*48*48
self.down_block2 = NetDownBlock(48, 64, 2) #24*24*24
self.down_block3 = NetDownBlock(64, 128, 2) #12*12*12
self.down_block4 = NetDownBlock(128, 256, 2) #6*6*6
self.up_block3 = NetUpBlock(256, 128, 256, 2) #12
self.up_block4 = NetUpBlock(256, 64, 160, 2) #24
self.out24_1 = NetInBlock(160, 64, 2)
self.out24_2 = NetInBlock(64, 32, 2)
self.out24_3 = NetOutSingleBlock(32, classes)
self.up_block5 = NetUpBlock(160, 48, 80, 2) #48
self.out48_1 = NetInBlock(80, 32, 2)
self.out48_2 = NetInBlock(64, 32, 2)
self.out48_3 = NetOutSingleBlock(32, classes)
self.up_block6 = NetJustUpBlock(80, 32, 2) #96
self.out96_1 = NetInBlock(64, 32, 1)
self.out96_block = NetOutSingleBlock(32, classes)
self.warp_layer96 = SpatialTransformer(96, 96, 96)
self.warp_layer48 = SpatialTransformer(48, 48, 48)
self.warp_layer24 = SpatialTransformer(24, 24, 24)
self.upsample = nn.Upsample(
scale_factor=2, mode='trilinear'
) #Interpolate(scale_factor=(2, 2, 2), mode='trilinear')
device = torch.device('cuda:0')
dir_deform = '.pth'
if __name__ == '__main__':
opt = TestOptions().parse()
deform_pred = Net_3d_ALL()
deform_pred.to(device)
deform_pred = torch.nn.DataParallel(deform_pred, [0])
tmp = torch.load(dir_deform)
deform_pred.load_state_dict(tmp)
deform_pred.eval()
warp_layer24 = SpatialTransformer(24,24,24)
warp_layer24 = warp_layer24.to(device)
warp_layer48 = SpatialTransformer(48,48,48)
warp_layer48 = warp_layer48.to(device)
warp_layer96 = SpatialTransformer(96,96,96)
warp_layer96 = warp_layer96.to(device)
model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = 0
total_valid = 0
train_loss = []
valid_loss = []
train_loss1 = []
if __name__ == '__main__':
opt = TrainOptions().parse() # Define the parameters
dataset_size = len(training_rows)
print('#training images = %d' % dataset_size)
STM_Net = Net_3d_ALL()
STM_Net.to(device)
STM_Net = torch.nn.DataParallel(STM_Net, [0])
tmp = torch.load(STM_Net_Dir)
STM_Net.load_state_dict(tmp)
STM_Net.eval()
warp_layer96 = SpatialTransformer(96,96,96)
warp_layer96 = warp_layer96.to(device)
model = create_model(opt)
visualizer = Visualizer(opt)
total_steps = 0
total_valid = 0
trainl1_96_loss = []
trainl1_48_loss = []
trainl1_24_loss = []
train_gradient = []
train_R_img = []
train_R_def = []
train_loss1 = []