def main():
json_path = os.path.join(args.model_dir)
params = utils.Params(json_path)
net = Unet(params.model).cuda()
# TODO - check exists
#checkpoint = torch.load('./final.pth.tar')
#net.load_state_dict(checkpoint)
train_dataset = AudioDataset(data_type='train')
test_dataset = AudioDataset(data_type='val')
train_data_loader = DataLoader(dataset=train_dataset,
batch_size=args.batch_size,
collate_fn=train_dataset.collate,
shuffle=True,
num_workers=4)
test_data_loader = DataLoader(dataset=test_dataset,
batch_size=args.batch_size,
collate_fn=test_dataset.collate,
shuffle=False,
num_workers=4)
torch.set_printoptions(precision=10, profile="full")
# Optimizer
optimizer = optim.Adam(net.parameters(), lr=1e-3)
# Learning rate scheduler
scheduler = ExponentialLR(optimizer, 0.95)
for epoch in range(args.num_epochs):
train_bar = tqdm(train_data_loader)
for input in train_bar:
train_mixed, train_clean, seq_len = map(lambda x: x.cuda(), input)
mixed = stft(train_mixed).unsqueeze(dim=1)
real, imag = mixed[..., 0], mixed[..., 1]
out_real, out_imag = net(real, imag)
out_real, out_imag = torch.squeeze(out_real,
1), torch.squeeze(out_imag, 1)
out_audio = istft(out_real, out_imag, train_mixed.size(1))
out_audio = torch.squeeze(out_audio, dim=1)
for i, l in enumerate(seq_len):
out_audio[i, l:] = 0
librosa.output.write_wav(
'mixed.wav', train_mixed[0].cpu().data.numpy()
[:seq_len[0].cpu().data.numpy()], 16000)
librosa.output.write_wav(
'clean.wav', train_clean[0].cpu().data.numpy()
[:seq_len[0].cpu().data.numpy()], 16000)
librosa.output.write_wav(
'out.wav', out_audio[0].cpu().data.numpy()
[:seq_len[0].cpu().data.numpy()], 16000)
loss = wSDRLoss(train_mixed, train_clean, out_audio)
print(epoch, loss)
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
torch.save(net.state_dict(), './final.pth.tar')
def main():
json_path = os.path.join(args.conf)
params = utils.Params(json_path)
net = Unet(params.model).cuda()
# TODO - check exists
# if os.path.exists('./ckpt/final.pth.tar'):
# checkpoint = torch.load('./ckpt/final.pth.tar')
# net.load_state_dict(checkpoint)
train_dataset = AudioDataset(data_type='train')
# test_dataset = AudioDataset(data_type='val')
train_data_loader = DataLoader(dataset=train_dataset, batch_size=args.batch_size,
collate_fn=train_dataset.collate, shuffle=True, num_workers=0)
# test_data_loader = DataLoader(dataset=test_dataset, batch_size=args.batch_size,
# collate_fn=test_dataset.collate, shuffle=False, num_workers=4)
# torch.set_printoptions(precision=10, profile="full")
# Optimizer
optimizer = optim.Adam(net.parameters(), lr=1e-2)
# Learning rate scheduler
scheduler = ExponentialLR(optimizer, 0.996)
if not os.path.exists('ckpt'): # model save dir
os.mkdir('ckpt')
for epoch in range(1, args.num_epochs+1):
train_bar = tqdm(train_data_loader, ncols=60)
loss_sum = 0.0
step_cnt = 0
for input_ in train_bar:
train_mixed, train_clean, seq_len = map(lambda x: x.cuda(), input_)
mixed = stft(train_mixed).unsqueeze(dim=1)
real, imag = mixed[..., 0], mixed[..., 1]
out_real, out_imag = net(real, imag)
out_real, out_imag = torch.squeeze(out_real, 1), torch.squeeze(out_imag, 1)
out_audio = istft(out_real, out_imag, train_mixed.size(1))
out_audio = torch.squeeze(out_audio, dim=1)
for i, l in enumerate(seq_len):
out_audio[i, l:] = 0
# librosa.output.write_wav('mixed.wav', train_mixed[0].cpu().data.numpy()[:seq_len[0].cpu().data.numpy()], 16000)
# librosa.output.write_wav('clean.wav', train_clean[0].cpu().data.numpy()[:seq_len[0].cpu().data.numpy()], 16000)
# librosa.output.write_wav('out.wav', out_audio[0].cpu().data.numpy()[:seq_len[0].cpu().data.numpy()], 16000)
loss = wSDRLoss(train_mixed, train_clean, out_audio)
# print(epoch, loss.item(), end='', flush=True)
loss_sum += loss.item()
step_cnt += 1
optimizer.zero_grad()
loss.backward()
optimizer.step()
avg_loss = loss_sum / step_cnt
print('epoch %d> Avg_loss: %.6f.\n' % (epoch, avg_loss))
scheduler.step()
if epoch %20 == 0:
torch.save(net.state_dict(), './ckpt/step%05d.pth.tar' % epoch)
def test(args):
# Setup Dataloader
data_json = json.load(open('config.json'))
data_path = data_json[args.dataset]['data_path']
t_loader = SaltLoader(data_path, split="test")
test_df=t_loader.test_df
test_loader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8)
# load Model
if args.arch=='unet':
model = Unet(start_fm=16)
else:
model=Unet_upsample(start_fm=16)
model_path = data_json[args.model]['model_path']
model.load_state_dict(torch.load(model_path)['model_state'])
model.cuda()
total = sum([param.nelement() for param in model.parameters()])
print('Number of params: %.2fM' % (total / 1e6))
#test
pred_list=[]
for images in test_loader:
images = Variable(images.cuda())
y_preds = model(images)
y_preds_shaped = y_preds.reshape(-1, args.img_size_target, args.img_size_target)
for idx in range(args.batch_size):
y_pred = y_preds_shaped[idx]
pred = torch.sigmoid(y_pred)
pred = pred.cpu().data.numpy()
pred_ori = resize(pred, (args.img_size_ori, args.img_size_ori), mode='constant', preserve_range=True)
pred_list.append(pred_ori)
#submit the test image predictions.
threshold_best=args.threshold
pred_dict = {idx: RLenc(np.round(pred_list[i] > threshold_best)) for i, idx in
enumerate(tqdm_notebook(test_df.index.values))}
sub = pd.DataFrame.from_dict(pred_dict, orient='index')
sub.index.names = ['id']
sub.columns = ['rle_mask']
sub.to_csv('./results/{}_submission.csv'.format(args.model))
print("The submission.csv saved in ./results")
def test(args):
# Setup Data
data_json = json.load(open('config.json'))
x = Variable(torch.randn(32, 1, 128, 128))
x = x.cuda()
# load Model
if args.arch == 'unet':
model = Unet(start_fm=16)
else:
model = Unet_upsample(start_fm=16)
model_path = data_json[args.model]['model_path']
model.load_state_dict(torch.load(model_path)['model_state'])
model.cuda()
total = sum([param.nelement() for param in model.parameters()])
print('Number of params: %.2fM' % (total / 1e6))
#visualize
y = model(x)
g = make_dot(y)
g.render('k')
weights_dict = torch.load(rootDir +
'/linear_factor=1_validation=6_test=7_unet3d.pt')
unet3d.load_state_dict(weights_dict)
resnet = ResBlock(
input_dim=2,
filter_dim=32,
output_dim=1,
)
resnet.to(device1)
weights_dict = torch.load(rootDir +
'/linear_factor=1_validation=6_test=7_resnet.pt')
resnet.load_state_dict(weights_dict)
# optimizer
optimizer0 = optim.Adam(unet3d.parameters(), lr=lr, betas=(0.5, 0.999))
optimizer1 = optim.Adam(resnet.parameters(), lr=lr, betas=(0.5, 0.999))
# dataloaders
# dataLoader_train_MS = Patient_data_loader_all(patientType='MS_old', flag_RDF_input=1)
# trainLoader_MS = data.DataLoader(dataLoader_train_MS, batch_size=batch_size, shuffle=True)
dataLoader_train_ICH = Patient_data_loader_all(patientType='ICH',
flag_RDF_input=1)
trainLoader_ICH = data.DataLoader(dataLoader_train_ICH,
batch_size=batch_size,
shuffle=True)
epoch = 0
gen_iterations = 1
display_iters = 5
'w': 'he_normal',
'b': 'normal'
})
criterion = nn.BCELoss(size_average=False)
else:
print("Choose a model.\nAborting....")
sys.exit()
if not os.path.exists('logs'):
os.mkdir('logs')
logFile = 'logs/' + fName + '.txt'
makeLogFile(logFile)
net.to(device)
optimizer = torch.optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-5)
nTrain = len(train_loader)
nValid = len(valid_loader)
nTest = len(test_loader)
minLoss = 1e8
convIter = 0
convCheck = 20
for epoch in range(args.epochs):
trLoss = []
vlLoss = []
vlGed = [0]
klEpoch = [0]
recEpoch = [0]
output_channels=1,
num_filters=[2**i for i in range(5, 10)], # or range(3, 8)
use_deconv=1,
flag_rsa=0)
resnet = ResBlock(
input_dim=2,
filter_dim=32,
output_dim=1,
)
unet3d.to(device)
resnet.to(device)
# optimizer
optimizer = optim.Adam(list(unet3d.parameters()) +
list(resnet.parameters()),
lr=lr,
betas=(0.5, 0.999))
ms = [0.3, 0.5, 0.7, 0.9]
ms = [np.floor(m * niter).astype(int) for m in ms]
scheduler = MultiStepLR(optimizer, milestones=ms, gamma=0.5)
# logger
logger = Logger('logs', rootDir, opt['linear_factor'],
opt['case_validation'], opt['case_test'])
# dataloader
dataLoader_train = COSMOS_data_loader(
split='Train',
patchSize=patchSize,
def validate(args):
# Setup Dataloader
data_json = json.load(open('config.json'))
data_path = data_json[args.dataset]['data_path']
v_loader = SaltLoader(data_path, split='val')
train_df = v_loader.train_df
val_loader = data.DataLoader(v_loader,
batch_size=args.batch_size,
num_workers=8)
# load Model
if args.arch == 'unet':
model = Unet(start_fm=16)
else:
model = Unet_upsample(start_fm=16)
model_path = data_json[args.model]['model_path']
model.load_state_dict(torch.load(model_path)['model_state'])
model.cuda()
total = sum([param.nelement() for param in model.parameters()])
print('Number of params: %.2fM' % (total / 1e6))
#validate
pred_list = []
for images, masks in val_loader:
images = Variable(images.cuda())
y_preds = model(images)
# print(y_preds.shape)
y_preds_shaped = y_preds.reshape(-1, args.img_size_target,
args.img_size_target)
for idx in range(args.batch_size):
y_pred = y_preds_shaped[idx]
pred = torch.sigmoid(y_pred)
pred = pred.cpu().data.numpy()
pred_ori = resize(pred, (args.img_size_ori, args.img_size_ori),
mode='constant',
preserve_range=True)
pred_list.append(pred_ori)
preds_valid = np.array(pred_list)
y_valid_ori = np.array(
[train_df.loc[idx].masks for idx in v_loader.ids_valid])
#jaccard score
accuracies_best = 0.0
for threshold in np.linspace(0, 1, 11):
ious = []
for y_pred, mask in zip(preds_valid, y_valid_ori):
prediction = (y_pred > threshold).astype(int)
iou = jaccard_similarity_score(mask.flatten(),
prediction.flatten())
ious.append(iou)
accuracies = [
np.mean(ious > iou_threshold)
for iou_threshold in np.linspace(0.5, 0.95, 10)
]
if accuracies_best < np.mean(accuracies):
accuracies_best = np.mean(accuracies)
threshold_best = threshold
print('Threshold: %.1f, Metric: %.3f' %
(threshold, np.mean(accuracies)))
print("jaccard score gets threshold_best=", threshold_best)
#other score way
thresholds = np.linspace(0, 1, 50)
ious = np.array([
iou_metric_batch(y_valid_ori, np.int32(preds_valid > threshold))
for threshold in tqdm_notebook(thresholds)
])
#don't understand
threshold_best_index = np.argmax(ious[9:-10]) + 9
iou_best = ious[threshold_best_index]
threshold_best = thresholds[threshold_best_index]
print("other way gets iou_best=", iou_best, "threshold_best=",
threshold_best)
def train(args):
# Setup Dataloader
data_json = json.load(open('config.json'))
data_path=data_json[args.dataset]['data_path']
t_loader = SaltLoader(data_path, img_size_ori=args.img_size_ori,img_size_target=args.img_size_target)
v_loader = SaltLoader(data_path, split='val',img_size_ori=args.img_size_ori, img_size_target=args.img_size_target)
train_loader = data.DataLoader(t_loader, batch_size=args.batch_size, num_workers=8, shuffle=True)
val_loader = data.DataLoader(v_loader, batch_size=args.batch_size, num_workers=8)
# Setup Model
if args.arch=='unet':
model = Unet(start_fm=16)
else:
model=Unet_upsample(start_fm=16)
print(model)
total = sum([param.nelement() for param in model.parameters()])
print('Number of params: %.2fM' % (total / 1e6))
model.cuda()
# Check if model has custom optimizer / loss
optimizer = torch.optim.Adam(model.parameters(), lr=args.l_rate)
loss_fn= nn.BCEWithLogitsLoss()
best_loss=100
mean_train_losses = []
mean_val_losses = []
for epoch in range(args.n_epoch):
train_losses = []
val_losses = []
for images, masks in train_loader:
images = Variable(images.cuda())
masks = Variable(masks.cuda())
outputs = model(images)
loss = loss_fn(outputs, masks)
train_losses.append(loss.data)
optimizer.zero_grad()
loss.backward()
optimizer.step()
for images, masks in val_loader:
images = Variable(images.cuda())
masks = Variable(masks.cuda())
outputs = model(images)
loss = loss_fn(outputs, masks)
val_losses.append(loss.data)
mean_train_losses.append(np.mean(train_losses))
mean_val_losses.append(np.mean(val_losses))
if np.mean(val_losses) < best_loss:
best_loss = np.mean(val_losses)
state = {'epoch': epoch + 1,
'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(), }
torch.save(state, "./saved_models/{}_{}_best_model.pkl".format(args.arch, args.dataset))
# Print Loss
print('Epoch: {}. Train Loss: {}. Val Loss: {}'.format(epoch + 1, np.mean(train_losses), np.mean(val_losses)))
state = {'model_state': model.state_dict(),
'optimizer_state': optimizer.state_dict(), }
torch.save(state, "./saved_models/{}_{}_final_model.pkl".format(args.arch, args.dataset))
print("saved two models in ./saved_models")
D_smv = np.real(S * D)
# network
model = Unet(input_channels=1,
output_channels=2,
num_filters=[2**i for i in range(5, 10)],
bilateral=1,
use_deconv=1,
use_deconv2=1,
renorm=0,
flag_r_train=0,
flag_UTFI=1)
model.to(device)
# optimizer
optimizer = optim.Adam(model.parameters(), lr=lr, betas=(0.5, 0.999))
Validation_loss = []
loss_total_list = []
while epoch < niter:
# training phase
model.train()
for idx, (ifreqs, masks, data_weights, wGs) in enumerate(trainLoader):
ifreqs = ifreqs.to(device, dtype=torch.float)
masks = masks.to(device, dtype=torch.float)
data_weights = data_weights.to(device, dtype=torch.float)
wGs = wGs.to(device, dtype=torch.float)
loss_PDF, loss_fidelity, loss_tv = utfi_train(model,
optimizer,
