def draw_sub(args, val_loader, model, epoch):
Channel_location = [
"FP1", "FP2", "F7", "F3", "FZ", "F4", "F8", "T7", "C3", "CZ", "C4",
"T8", "P7", "P3", "PZ", "P4", "P8", "O1", "O2"
]
model.eval()
for i, (input, target, max_num) in enumerate(val_loader):
if args.onGPU == True:
input = input.cuda()
target = target.cuda()
with torch.no_grad():
# run the mdoel
decode = model(input)
batch, channel = 0, 14
print("decode(shape): ", decode.shape)
for j in range(19):
filename = Channel_location[j] + 'e{}'.format(str(str(epoch)))
snr_i_t = numpy_SNR(input[batch][j], target[batch][j])
snr_i_d = numpy_SNR(input[batch][j], decode[batch][j])
draw_psd(filename, input[batch][j], 256, snr_i_t, snr_i_d, 1, 2, 1)
draw_psd(filename, decode[batch][j], 256, snr_i_t, snr_i_d, 1, 2,
2)
draw_psd(filename, target[batch][j], 256, snr_i_t, snr_i_d, 1, 2,
3)
imgSave(args.savefig + "/sub/", str(j) + "_" + filename)
loss4one = []
loss4sec = []
loss4third = []
criterion = nn.MSELoss()
for j in range(int(decode.shape[1])):
loss = criterion(decode[:, j, :], target[:, j, :])
doutput = decode[:, j, 1:] - decode[:, j, :-1]
dtarget = target[:, j, 1:] - target[:, j, :-1]
loss2 = criterion(doutput, dtarget)
# print("train(shape):", doutput.shape)
d2output = doutput[:, 1:] - doutput[:, :-1]
d2target = dtarget[:, 1:] - dtarget[:, :-1]
loss3 = criterion(d2output, d2target)
loss4one.append(loss.item())
loss4sec.append(loss2.item())
loss4third.append(loss3.item())
fp = open(args.savefig + "/sub/" + "Channel_loss.txt", "a")
fp.write("MSE(one): " + str(loss4one))
fp.write("\nMSE(sec): " + str(loss4sec))
fp.write("\nMSE(Third): " + str(loss4third))
fp.close()
print("MSE(one): ", loss4one)
print("MSE(sec): ", loss4sec)
print("MSE(Third): ", loss4third)
def draw(args, val_loader, model, epoch):
'''
:param args: general arguments
:param val_loader: loaded for validation dataset
:param model: model
:return: non
'''
Channel_location = [
"FP1", "FP2", "F7", "F3", "FZ", "F4", "F8", "FT7", "FC3", "FCZ", "FC4",
"FT8", "T4", "C3", "FZ", "C4", "T4", "TP7", "CP3", "CPZ", "CP4", "TP8",
"T5", "P3", "PZ", "P4", "T6", "O1", "OZ", "O2"
]
model.eval()
for i, (input, target, max_num) in enumerate(val_loader):
if args.onGPU == True:
input = input.cuda()
target = target.cuda()
with torch.no_grad():
# run the mdoel
decode = model(input)
batch, channel = 0, 14
#print("draw(max): ", max_num[batch])
print("draw(shape): ", input.shape)
target = target * max_num[batch]
snr = numpy_SNR(target[batch][channel], target[batch][channel])
draw_raw("target", target[batch][channel], 1, snr)
input = input * max_num[batch]
snr = numpy_SNR(input[batch][channel], target[batch][channel])
draw_raw("input", input[batch][channel], 2, snr)
decode = decode * max_num[batch]
snr = numpy_SNR(decode[batch][channel], target[batch][channel])
draw_raw("decode", decode[batch][channel], 3, snr)
#imgSave(epoch)
snr_i_t = numpy_SNR(input[batch][channel], target[batch][channel])
snr_i_d = numpy_SNR(input[batch][channel], decode[batch][channel])
draw_psd(str(epoch), input[batch][channel], 256, snr_i_t, snr_i_d, 1,
2, 1)
draw_psd(str(epoch), decode[batch][channel], 256, snr_i_t, snr_i_d, 1,
2, 2)
draw_psd(str(epoch), target[batch][channel], 256, snr_i_t, snr_i_d, 1,
2, 3)
imgSave(args.savefig, '/e{}'.format(str(str(epoch) + "PSD")))
break