Exemplo n.º 1
0
# Save all the codes
os.system('mkdir %s' % opt.experiment )
os.system('cp *.py %s' % opt.experiment )

if torch.cuda.is_available() and opt.noCuda:
    print("WARNING: You have a CUDA device, so you should probably run with --cuda")

# Initialize image batch
imBatch = Variable(torch.FloatTensor(opt.batchSize, 3, 300, 300) )
labelBatch = Variable(torch.FloatTensor(opt.batchSize, opt.numClasses, 300, 300) )
maskBatch = Variable(torch.FloatTensor(opt.batchSize, 1, 300, 300) )
labelIndexBatch = Variable(torch.LongTensor(opt.batchSize, 1, 300, 300) )

# Initialize network
if opt.isDilation:
    encoder = model.encoderDilation()
    decoder = model.decoderDilation()
elif opt.isSpp:
    encoder = model.encoderSPP()
    decoder = model.decoderSPP()
else:
    encoder = model.encoder()
    decoder = model.decoder()

encoder.load_state_dict(torch.load('%s/encoder_%d.pth' % (opt.modelRoot, opt.epochId) ) )
decoder.load_state_dict(torch.load('%s/decoder_%d.pth' % (opt.modelRoot, opt.epochId) ) )
encoder = encoder.eval()
decoder = decoder.eval()

# Move network and containers to gpu
if not opt.noCuda:
maskBatch = Variable(torch.FloatTensor(opt.batchSize, 1, 300, 300))
labelIndexBatch = Variable(torch.LongTensor(opt.batchSize, 1, 300, 300))

# Initialize network
encoder_normal = model.encoder()
decoder_normal = model.decoder()
model_root_normal = '/datasets/cse152-252-sp20-public/unet_checkpoints/unet_original_zq'
epoch_id_normal = 181
encoder_normal.load_state_dict(
    torch.load('%s/encoder_%d.pth' % (model_root_normal, epoch_id_normal)))
decoder_normal.load_state_dict(
    torch.load('%s/decoder_%d.pth' % (model_root_normal, epoch_id_normal)))
encoder_normal = encoder_normal.eval()
decoder_normal = decoder_normal.eval()

encoder_dilation = model.encoderDilation()
decoder_dilation = model.decoderDilation()
model_root_dilation = '/datasets/cse152-252-sp20-public/unet_checkpoints/unet_original_zq_dilation'
epoch_id_dilation = 180
encoder_dilation.load_state_dict(
    torch.load('%s/encoder_%d.pth' % (model_root_dilation, epoch_id_dilation)))
decoder_dilation.load_state_dict(
    torch.load('%s/decoder_%d.pth' % (model_root_dilation, epoch_id_dilation)))
encoder_dilation = encoder_dilation.eval()
decoder_dilation = decoder_dilation.eval()

encoder_spp = model.encoderDilation()
decoder_spp = model.decoderDilation(isSpp=True)
model_root_spp = '/datasets/cse152-252-sp20-public/unet_checkpoints/unet_original_zq_spp'
epoch_id_spp = 180
encoder_spp.load_state_dict(