def test_model(batch):
batch_size, n_channels, imsize, imsize = batch.shape
model = UNet()
assert model(batch).shape == (batch_size, 1, imsize, imsize)
def main_worker(train_loader, val_loader, args):
global best_loss
# create model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f'=> device used: {device}')
norm_kwargs = {
'mode': args.norm_mode,
'alpha_fwd': args.afwd,
'alpha_bkw': args.abkw,
'ecm': args.ecm
}
print("=> creating model...")
model = UNet(args.classes, norm_layer=norm_layer,
norm_kwargs=norm_kwargs).to(device)
print(model)
print("=> creating optimizer...")
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
print("=> setting up learning rate scheduler...")
scheduler = lr_scheduler.StepLR(optimizer,
step_size=args.lr_milestone,
gamma=args.lr_multiplier)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_loss = checkpoint['best_loss']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
scheduler.load_state_dict(checkpoint['scheduler'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
cudnn.benchmark = False if args.seed else True
if args.evaluate:
validate(val_loader, model, args.start_epoch, device, args)
return
for epoch in range(args.start_epoch, args.epochs):
if epoch: scheduler.step()
# train for one epoch
train(train_loader, model, optimizer, epoch, device, args)
# evaluate on validation set
eval_loss = validate(val_loader, model, epoch, device, args)
# remember best loss and save checkpoint
is_best = eval_loss < best_loss
best_loss = min(eval_loss, best_loss)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_loss': best_loss,
'optimizer': optimizer.state_dict(),
'scheduler': scheduler.state_dict(),
}, is_best, args)
print('best val loss: {:4f}'.format(best_loss))
# load best model weights
model_best_file = os.path.join(args.model_dir, 'model_best.pth.tar')
if os.path.isfile(model_best_file):
print("=> loading checkpoint '{}'".format(model_best_file))
checkpoint = torch.load(model_best_file)
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
model_best_file, checkpoint['epoch']))
return model
def model_denoising(filename, model_type='Unet'):
audio = audio_files_to_numpy(UPLOAD_FOLDER, [filename], SAMPLE_RATE,
FRAME_LENGTH, HOP_LENGTH_FRAME, MIN_DURATION)
# Squared spectrogram dimensions
dim_square_spec = int(N_FFT / 2) + 1
m_amp_db, m_pha = numpy_audio_to_matrix_spectrogram(
audio, dim_square_spec, N_FFT, HOP_LENGTH_FFT)
# get device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# transform
X_in = torch.from_numpy(scaled_in(m_amp_db)).unsqueeze(1).to(
device, dtype=torch.float)
if model_type == 'Unet':
model = UNet(start_fm=32).to(device)
try:
if torch.cuda.is_available():
model.load_state_dict(torch.load('./model/unet.pth'))
else:
model.load_state_dict(
torch.load('./model/unet.pth', map_location='cpu'))
except:
st.error(
'Your weight are not found. Make sure the weight located in ./model'
)
else:
model = UNet_ResNet(start_fm=16).to(device)
try:
if torch.cuda.is_available():
model.load_state_dict(torch.load('./model/unetres.pth'))
else:
model.load_state_dict(
torch.load('./model/unetres.pth', map_location='cpu'))
except:
st.error(
'Your weight are not found. Make sure the weight located in ./model'
)
with torch.no_grad():
X_pred = model(X_in)
pred_amp_db = inv_scaled_ou(X_pred.squeeze().detach().cpu().numpy(),
REDUCE_RATE)
X_denoise = m_amp_db - pred_amp_db
ou_audio = matrix_spectrogram_to_numpy_audio(X_denoise, m_pha,
FRAME_LENGTH, HOP_LENGTH_FFT)
nb_samples = ou_audio.shape[0]
denoise_long = ou_audio.reshape(1, nb_samples * FRAME_LENGTH) * 10
sf.write(UPLOAD_FOLDER + 'out_' + filename, denoise_long[0, :],
SAMPLE_RATE)
return m_amp_db, m_pha, pred_amp_db, X_denoise
validation_dataset = LTSIDDataset(cfg.input_dir, cfg.truth_dir, preprocess=cfg.preprocess,
preprocess_dir=cfg.preprocess_dir, collection='validation',
transforms=transforms.Compose([
trf.RandomCrop(cfg.patch_size),
trf.ToTensor(),
trf.RandomHorizontalFlip(p=0.5),
trf.RandomVerticalFlip(p=0.5),
trf.RandomTranspose(p=0.5),
]))
train_loader = DataLoader(train_dataset, batch_size=cfg.batch_size, shuffle=True)
validation_loader = DataLoader(validation_dataset, batch_size=cfg.batch_size, shuffle=True,)
print('Dataset loaded!')
# Set up model
model = UNet().to(device)
# Set up loss function
loss_func = nn.L1Loss()
# loss_func = perceptual_loss(perceptual_model='vgg16', dist_func=nn.MSELoss(), device=device) # Perceptual loss
# Set up optimizer
optimizer = optim.Adam(model.parameters(), lr=cfg.initial_learning_rate)
# Experiment with 16-bit precision
amp.initialize(model, optimizer, opt_level='O2')
# Learning rate scheduling
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=cfg.epochs/2, gamma=0.1)
except:
artifact = run.use_artifact('Spectrograms:latest', type='Dataset')
artifact_dir = artifact.download(DATA_PATH)
# load dataset
dataset = SpeechDataset(DATA_PATH)
trainloader, validloader = get_dataloader(dataset=dataset,
batch_size=args.batchsize)
# get model and define loss func, optimizer
n_classes = N_CLASSES
epochs = args.epoch
tag = 'Unet'
if args.model == 'Unet':
model = UNet(start_fm=args.startfm).to(device)
else:
tag = 'UnetRes'
model = UNet_ResNet(dropout=args.dropout,
start_fm=args.startfm).to(device)
run.tags = [tag]
criterion = nn.SmoothL1Loss()
# loss_func = Weighted_Cross_Entropy_Loss()
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# wandb watch
run.watch(models=model, criterion=criterion, log='all', log_freq=10)
1))),
normalize=True,
mean=mean,
std=std)
data_testB = DataLoader(testB, batch_size=batch_size, shuffle=False)
#set device (eiter cpu or gpu)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)
use = 'segnet'
if use == 'segnet':
model = SegNet().to(device)
summary(model, (3, 256, 256))
else:
model = UNet().to(device)
summary(model, (3, 256, 256))
F1_scores = train(model, optim.AdamW(model.parameters(), lr=0.0001), 100,
data_tr, data_val, data_testA, data_testB, device)
plt.rcParams['figure.figsize'] = [18, 12]
ig, (ax0, ax1, ax2, ax3) = plt.subplots(nrows=4, sharex=True)
ax0.errorbar(range(len(F1_scores['train_mean'])),
F1_scores['train_mean'],
yerr=F1_scores['train_std'],
fmt='-o')
ax0.set_title('F1-Score train')
ax1.errorbar(range(len(F1_scores['val_mean'])),
with open(config_file, "r") as ymlfile:
yml_file = yaml.load(ymlfile)
cfg = Config(yml_file)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = LTSIDDataset(cfg.input_dir,
cfg.truth_dir,
preprocess=cfg.preprocess,
preprocess_dir=cfg.preprocess_dir,
collection='test',
transforms=transforms.Compose([trf.ToTensor()]))
dataloader = DataLoader(dataset, batch_size=cfg.batch_size, shuffle=False)
# Load model
model = UNet().to(device)
model.eval() #Set model to evaluation mode
# Set up loss function
loss_func = nn.L1Loss()
checkpoint = torch.load(cfg.checkpoint_to_load)
model.load_state_dict(checkpoint['state_dict'])
# Make model checkpoint dir
Path(cfg.output_dir).mkdir(exist_ok=True)
counter = 0
with torch.no_grad():
test_loss = 0.0
test_len = len(dataloader)