├── .gitignore ├── input └── air.npy ├── requirements.txt ├── result └── air_recover.npy ├── data.py ├── impute.py ├── utils.py ├── train.py └── TCDAE.py /.gitignore: -------------------------------------------------------------------------------- 1 | /__pycache__ 2 | /output 3 | -------------------------------------------------------------------------------- /input/air.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yeguo419/DAE_impute/HEAD/input/air.npy -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | torch==1.0.1.post2 2 | numpy==1.18.0 3 | matplotlib==3.0.3 4 | -------------------------------------------------------------------------------- /result/air_recover.npy: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/yeguo419/DAE_impute/HEAD/result/air_recover.npy -------------------------------------------------------------------------------- /data.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.utils.data as data 3 | import numpy as np 4 | 5 | 6 | def TDMinMaxScaler(data): 7 | #mi = data[np.where(data != np.nan)].min(axis=1).min(axis=0) 8 | #ma = data[np.where(data != np.nan)].max(axis=1).max(axis=0) 9 | mi = data.min(axis=2) 10 | ma = data.max(axis=2) 11 | #print(mi.shape) 12 | 13 | 14 | 15 | for i in range(data.shape[0]): 16 | for j in range(data.shape[1]): 17 | data[i][j] = (data[i][j] - mi[i][j]) / (ma[i][j] - mi[i][j] + 1e-8) 18 | #data = (data - mi) / (ma - mi) 19 | return data, ma, mi 20 | 21 | class MissDataset(data.Dataset):# return (complete_data, missing_data, mask) 22 | 23 | def __init__(self, data , mask): 24 | super(MissDataset, self).__init__() 25 | 26 | self.data = torch.Tensor(data) 27 | 28 | self.mask = torch.Tensor(mask) 29 | print('Create Dataset, size:', data.shape) 30 | print('Sample, Var, Time') 31 | 32 | 33 | 34 | def __getitem__(self, index): 35 | return self.data[index], self.mask[index] 36 | 37 | def __len__(self): 38 | return len(self.data) 39 | 40 | 41 | def MissDataLoader(file_path): 42 | data = np.load(file_path) 43 | #print('Data loading, shape :', data.shape) 44 | data_m = np.ones(data.shape) 45 | data_m[np.where(np.isnan(data))] = 0 46 | 47 | data_x = data[:] 48 | data_x[np.where(np.isnan(data))] = 0 49 | 50 | data_x, ma, mi = TDMinMaxScaler(data_x) 51 | 52 | return data_x, data_m, ma, mi 53 | 54 | -------------------------------------------------------------------------------- /impute.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import torch 4 | from data import MissDataset, MissDataLoader 5 | from TCDAE import TCDAE 6 | from utils import * 7 | 8 | parser = argparse.ArgumentParser('Impute missing data Using TCDAE') 9 | 10 | parser.add_argument('--model_path', type=str, 11 | default='./output/final.pth', #required=True, 12 | help='Path to model file') 13 | 14 | parser.add_argument('--dataset_path', type=str, 15 | default='./input/air.npy' , 16 | help='File path of dataset to impute') 17 | 18 | parser.add_argument('--out_dir', type=str, 19 | default='./result', 20 | help='Directory putting imputed dataset') 21 | 22 | parser.add_argument('--use_cuda', type=int, default=0, 23 | help='Whether use GPU to impute data') 24 | 25 | 26 | def impute(args, model, z, m): 27 | 28 | tmp_x = model(torch.FloatTensor(z).cuda()) if args.use_cuda else model(torch.FloatTensor(z)) 29 | if args.use_cuda: 30 | tmp_x = tmp_x.cpu() 31 | 32 | tmp_x = tmp_x.detach().numpy() 33 | 34 | x = z * m + tmp_x * (1 - m) 35 | return x 36 | 37 | def re_scale(data, ma, mi): 38 | for i in range(data.shape[0]): 39 | for j in range(data.shape[1]): 40 | data[i][j] = data[i][j] * (ma[i][j] - mi[i][j] + 1e-8) + mi[i][j] 41 | return data 42 | 43 | def save_data(args, data): 44 | if not os.path.exists(args.out_dir): 45 | os.makedirs(args.out_dir) 46 | 47 | file_name = os.path.split(args.dataset_path)[-1].split('.')[0] + '_recover.npy' 48 | 49 | save_path = os.path.join(args.out_dir, file_name) 50 | 51 | np.save(save_path, data) 52 | 53 | if __name__ == '__main__': 54 | args = parser.parse_args() 55 | 56 | data_z, data_m, ma, mi = MissDataLoader(args.dataset_path) 57 | model = torch.load(args.model_path, map_location=lambda storage, loc: storage) 58 | 59 | model.eval() 60 | 61 | if args.use_cuda: 62 | model.cuda() 63 | 64 | data_r = impute(args, model, data_z, data_m) 65 | 66 | 67 | imputed_data = re_scale(data_r, ma, mi) 68 | 69 | save_data(args, imputed_data) 70 | -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | #encoding:utf-8 2 | import torch 3 | import os 4 | import torch.nn as nn 5 | import torch.nn.utils.rnn as rnn 6 | import torch.utils.data as data 7 | import torch.nn.functional as F 8 | from torch.autograd import Variable 9 | from math import isnan 10 | import numpy as np 11 | import matplotlib 12 | matplotlib.use('Agg') 13 | import matplotlib.pyplot as plt 14 | 15 | class MaskLoss(nn.Module): # loss function with mask(1:base_loss, 0:no_loss) 16 | def __init__(self, base_loss = nn.MSELoss(reduction='none')): 17 | super(MaskLoss, self).__init__() 18 | self.base_loss = base_loss 19 | 20 | def forward(self, output, target, a_m): 21 | a = 1 22 | #loss = torch.sum(self.base_loss(output, target) * mask) 23 | loss = torch.sqrt(torch.mean(torch.pow(a_m * output - a_m * target, 2))) 24 | #loss = torch.sqrt(torch.mean(torch.pow(a_m * output - a_m * target, 2) * 1 + (1 - a) * torch.pow((1 - a_m) * output - (1 - a_m)*target, 2)) ) 25 | return loss 26 | 27 | 28 | 29 | 30 | class MaskDataset(data.Dataset):#dataset with mask, return (data, mask) 31 | 32 | def __init__(self, data, mask): 33 | super(MaskDataset, self).__init__() 34 | self.data = torch.Tensor(data) 35 | self.mask = torch.Tensor(mask) 36 | print('Create MaskDataset, size:', data.shape) 37 | 38 | 39 | def __getitem__(self, index): 40 | return self.data[index], self.mask[index] 41 | 42 | def __len__(self): 43 | return len(self.data) 44 | 45 | 46 | class MissMaskDataset(data.Dataset):#dataset with mask and ramdom missing, return (incomplete_data, complete_data, mask) 47 | 48 | def __init__(self, data, mask, missing_rate): 49 | 50 | def set0(data, rate): 51 | m = np.random.uniform(0, 1, data.shape) 52 | n_data = np.copy(data) 53 | n_data[np.where(m < missing_rate)] = 0 54 | return n_data 55 | 56 | super(MissMaskDataset, self).__init__() 57 | self.missing_data = torch.Tensor(set0(data, missing_rate)) 58 | self.mask = torch.Tensor(mask) 59 | self.raw_data = torch.Tensor(data) 60 | 61 | print('Create MissMaskDataset, size:', data.shape) 62 | 63 | def __getitem__(self, index): 64 | return self.missing_data[index], self.raw_data[index], self.mask[index] 65 | 66 | def __len__(self): 67 | return len(self.raw_data) 68 | 69 | def draw_curve(loss_list, label, name): 70 | plt.xlabel('epoch') 71 | plt.ylabel(label) 72 | plt.plot(loss_list) 73 | plt.savefig(name, dpi=500) 74 | 75 | 76 | def cat_z_m(x, m): 77 | x = torch.cat((x.view(x.shape[0], x.shape[1], x.shape[2], 1), m.view(m.shape[0], m.shape[1], m.shape[2], 1)), 3).view(x.shape[0], x.shape[1], -1) 78 | return x 79 | 80 | 81 | def num_param(model, comment=''): 82 | parameters = filter(lambda p: p.requires_grad, model.parameters()) 83 | parameters = sum([np.prod(p.size()) for p in parameters]) 84 | print(comment + ' Trainable Parameters: %.4f million' % (parameters / 1000000)) 85 | print(comment + ' Usage is %.4f' % (parameters * 32 / 8 / 1024 / 1024)) 86 | 87 | 88 | def get_step_loss(model, z, m): 89 | r = model(torch.Tensor(z), torch.Tensor(m)).cpu().detach().numpy() 90 | loss = np.sqrt(np.mean(np.square(m * r - m * z))) 91 | return loss -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import os 3 | import torch 4 | from data import MissDataset, MissDataLoader 5 | from TCDAE import TCDAE 6 | from utils import * 7 | 8 | import random 9 | import math 10 | import torch.nn as nn 11 | import torch.nn.utils.rnn as rnn 12 | import torch.utils.data as data 13 | import torch.nn.functional as F 14 | from torch.autograd import Variable 15 | import numpy as np 16 | import matplotlib 17 | matplotlib.use('Agg') 18 | import matplotlib.pyplot as plt 19 | 20 | parser = argparse.ArgumentParser( 21 | "Time series missing data imputation" 22 | "with Temporal Convolutional Denoising Autoencoder") 23 | 24 | parser.add_argument('--train_path', type=str, default='./input/air.npy', 25 | help='directory of train data, with npy format') 26 | 27 | 28 | parser.add_argument('--B', default=128, type=int, 29 | help='Number of channels in bottleneck 1 × 1-conv block') 30 | parser.add_argument('--H', default=128, type=int, 31 | help='Number of channels in convolutional blocks') 32 | parser.add_argument('--P', default=3 , type=int, 33 | help='Kernel size in convolutional blocks') 34 | parser.add_argument('--L', default=2, type=int, 35 | help='Number of convolutional blocks in each repeat') 36 | parser.add_argument('--R', default=2, type=int, 37 | help='Number of repeats') 38 | parser.add_argument('--C', default=0.2, type=float, 39 | help='Probability of dropout in input layer') 40 | 41 | parser.add_argument('--epochs', default=10, type=int, 42 | help='Number of maximum epochs') 43 | parser.add_argument('--lr', default=1e-4, type=float, 44 | help='Init learning rate') 45 | parser.add_argument('--l2', default=1e-3, type=float, 46 | help='weight decay (L2 penalty)') 47 | 48 | # minibatch 49 | parser.add_argument('--shuffle', default=1, type=int, 50 | help='reshuffle the data at every epoch') 51 | parser.add_argument('--batch_size', default=64, type=int, 52 | help='Batch size') 53 | #parser.add_argument('--num_workers', default=4, type=int, 54 | # help='Number of workers to generate minibatch') 55 | 56 | parser.add_argument('--save_folder', default='./output/', 57 | help='Location to save epoch models') 58 | 59 | 60 | def train(dataLoader, args, train_z, train_m): 61 | 62 | 63 | model = TCDAE(N=args.N, B=args.B, H=args.H, P=args.P, X=args.L, R=args.R, C=args.C) 64 | num_param(model) 65 | model.cuda() 66 | model = torch.nn.DataParallel(model) 67 | 68 | 69 | optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.l2) 70 | 71 | loss_fn = MaskLoss() 72 | 73 | loss_list = [] 74 | best_rmse = 1e8 75 | 76 | for epoch in range(args.epochs): 77 | 78 | model.train() 79 | 80 | for i, (z, m) in enumerate(dataLoader): 81 | z = Variable(z).cuda() 82 | m = Variable(m).cuda() 83 | 84 | decoded = model(z, m) 85 | 86 | loss = loss_fn(decoded, z, m)#Denosing: calculate loss by decode and complete data 87 | #loss = loss_fn(decoded, raw_x) 88 | 89 | optimizer.zero_grad() 90 | loss.backward() 91 | optimizer.step() 92 | 93 | 94 | model.eval() 95 | 96 | step_loss = get_step_loss(model, train_z, train_m) 97 | 98 | 99 | loss_list.append(step_loss) 100 | 101 | #test_rmse = get_rmse(model, train_x, train_z, train_m) 102 | print (epoch, 'Step loss:', step_loss)#, ' Test rmse:', test_rmse) 103 | #if best_rmse > test_rmse: 104 | # best_rmse = test_rmse 105 | 106 | save_path = os.path.join(args.save_folder, 'model_epoch_%d.pth' % (epoch)) 107 | torch.save(model.module, save_path) 108 | #save(model, train_x, train_z, train_m) 109 | 110 | 111 | 112 | #draw_curve(loss_list, 'RMSE', 'RMSE') 113 | #print('Best Test RMSE', best_rmse) 114 | return model 115 | 116 | def main(args): 117 | train_z, train_m, _, _ = MissDataLoader(args.train_path) 118 | train_set = MissDataset(train_z, train_m) 119 | 120 | args.N = train_z.shape[1] 121 | train_loader = data.DataLoader(train_set, batch_size=args.batch_size, shuffle=args.shuffle) 122 | 123 | if not os.path.exists(args.save_folder): 124 | os.makedirs(args.save_folder) 125 | 126 | model = train(train_loader, args, train_z, train_m) 127 | torch.save(model.module, os.path.join(args.save_folder, 'final.pth')) 128 | 129 | if __name__ == '__main__': 130 | args = parser.parse_args() 131 | print(args) 132 | main(args) 133 | 134 | -------------------------------------------------------------------------------- /TCDAE.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | import torch.nn.functional as F 4 | EPS = 1e-8 5 | 6 | class TCDAE(nn.Module): 7 | def __init__(self, N, B=128, H=128, P=3, X=4, R=2, C=0.2): 8 | super(TCDAE, self).__init__() 9 | print(B, H, P, X, R, C) 10 | 11 | #B = N 12 | self.C = C 13 | 14 | self.encoder = Encoder(N, B) 15 | self.impute_net = TemporalConvNet(B, H, P, X, R) 16 | self.decoder = Decoder(N, B) 17 | 18 | self.tcnn = nn.Sequential( 19 | self.encoder, 20 | self.impute_net, 21 | self.decoder) 22 | 23 | def forward(self, x, m=None): 24 | #x M * N * T 25 | x = nn.Dropout(self.C)(x) 26 | 27 | 28 | return self.tcnn(x)#M * N * T 29 | 30 | 31 | 32 | class Encoder(nn.Module): 33 | def __init__(self, N, B): 34 | super(Encoder, self).__init__() 35 | 36 | self.encode_net = nn.Sequential( 37 | nn.Conv1d(N, B, 15, padding=7), 38 | nn.ReLU()) 39 | 40 | 41 | def forward(self, x): 42 | v = self.encode_net(x) 43 | return v 44 | 45 | class Decoder(nn.Module): 46 | def __init__(self, N, B): 47 | super(Decoder, self).__init__() 48 | 49 | self.decode_net = nn.Sequential( 50 | nn.Conv1d(B, N, 1), 51 | nn.Sigmoid()) 52 | 53 | def forward(self, x): 54 | return self.decode_net(x) 55 | 56 | 57 | 58 | class TemporalConvNet(nn.Module): 59 | def __init__(self, B, H, P, X, R, norm_type="gLN", causal=False, 60 | mask_nonlinear='relu'): 61 | """ 62 | 63 | mask_nonlinear: use which non-linear function to generate mask 64 | """ 65 | super(TemporalConvNet, self).__init__() 66 | # Hyper-parameter 67 | self.mask_nonlinear = mask_nonlinear 68 | # Components 69 | # [M, B, T] -> [M, B, T] 70 | layer_norm = ChannelwiseLayerNorm(B) 71 | 72 | repeats = [] 73 | for r in range(R): 74 | blocks = [] 75 | for x in range(X): 76 | dilation = 2**x 77 | padding = (P - 1) * dilation if causal else (P - 1) * dilation // 2 78 | blocks += [TemporalBlock(B, H, P, stride=1, 79 | padding=padding, 80 | dilation=dilation, 81 | norm_type=norm_type, 82 | causal=causal)] 83 | repeats += [nn.Sequential(*blocks)] 84 | temporal_conv_net = nn.Sequential(*repeats) 85 | #self.repeats = repeats 86 | # [M, B, K] -> [M, N, K] 87 | 88 | # Put together 89 | self.network = nn.Sequential(layer_norm, 90 | #bottleneck_conv1x1, 91 | temporal_conv_net) 92 | # nn.PReLU(), 93 | # mask_conv1x1) 94 | #nn.Sigmoid()) 95 | 96 | # mask_conv1x1 = nn.Conv1d(B, N, 1, bias=False) 97 | # self.mask_net = nn.Sequential(nn.PReLU(), 98 | # mask_conv1x1, 99 | # nn.Sigmoid()) 100 | 101 | def forward(self, x): 102 | """ 103 | x: M * B * T 104 | output M * B * T 105 | """ 106 | _, output = self.network(x) 107 | return output 108 | 109 | 110 | class TemporalBlock(nn.Module): 111 | def __init__(self, in_channels, out_channels, kernel_size, 112 | stride, padding, dilation, norm_type="gLN", causal=False): 113 | super(TemporalBlock, self).__init__() 114 | # [M, B, K] -> [M, H, K] 115 | conv1x1 = nn.Conv1d(in_channels, out_channels, 3, padding=1, bias=False) 116 | prelu = nn.PReLU() 117 | norm = chose_norm(norm_type, out_channels) 118 | # [M, H, K] -> [M, B, K] 119 | dsconv = DepthwiseSeparableConv(out_channels, in_channels, kernel_size, 120 | stride, padding, dilation, norm_type, 121 | causal) 122 | # Put together 123 | self.net = nn.Sequential(conv1x1, prelu, norm, dsconv) 124 | 125 | #self.att_net = nn.Sequential(nn.Conv1d(in_channels, 1, 3, padding=1), 126 | # nn.PReLU()) # 1 * K 127 | # self.att_net = nn.Sequential(nn.Conv1d(in_channels, in_channels, 3, padding=1), 128 | # nn.Sigmoid()) 129 | 130 | def forward(self, x): 131 | """ 132 | Args: 133 | x: [M, B, K], sc:[M, Sc, K] we simply set Sc == B 134 | Returns: 135 | [M, B, K] 136 | """ 137 | if type(x) != tuple: 138 | sc = torch.zeros(x.shape, dtype=torch.float32) 139 | if x.is_cuda: 140 | sc = sc.cuda() 141 | else: 142 | sc = x[1] 143 | x = x[0] 144 | 145 | residual = x 146 | #out = self.net(x) 147 | out, new_sc = self.net(x) 148 | 149 | #a = self.att_net(out) # [M, B, K] 150 | #a = torch.mean(a, dim=2, keepdim=True)# [M, B, 1] 151 | 152 | return out + residual, sc + new_sc# * a # look like w/o F.relu is better than w/ F.relu 153 | # return F.relu(out + residual) 154 | 155 | 156 | class DepthwiseSeparableConv(nn.Module): 157 | def __init__(self, in_channels, out_channels, kernel_size, 158 | stride, padding, dilation, norm_type="gLN", causal=False): 159 | super(DepthwiseSeparableConv, self).__init__() 160 | # Use `groups` option to implement depthwise convolution 161 | # [M, H, K] -> [M, H, K] 162 | depthwise_conv = nn.Conv1d(in_channels, in_channels, kernel_size, 163 | stride=stride, padding=padding, 164 | dilation=dilation, groups=in_channels, 165 | bias=False) 166 | if causal: 167 | chomp = Chomp1d(padding) 168 | prelu = nn.PReLU() 169 | norm = chose_norm(norm_type, in_channels) 170 | # [M, H, K] -> [M, B, K] 171 | #output = nn.Conv1d(in_channels, out_channels, 1, bias=False) 172 | self.output_conv = nn.Conv1d(in_channels, out_channels, 3, padding=1, bias=False) 173 | 174 | self.sc_conv = nn.Conv1d(in_channels, out_channels, 3, padding=1, bias=False) 175 | 176 | # Put together 177 | if causal: 178 | self.net = nn.Sequential(depthwise_conv, chomp, prelu, norm)#, output) 179 | else: 180 | self.net = nn.Sequential(depthwise_conv, prelu, norm)#, output) 181 | 182 | def forward(self, x): 183 | """ 184 | Args: 185 | x: [M, H, K] 186 | Returns: 187 | result: [M, B, K] 188 | """ 189 | x = self.net(x) 190 | output = self.output_conv(x) 191 | sc = self.sc_conv(x) 192 | 193 | 194 | return output, sc 195 | 196 | def chose_norm(norm_type, channel_size): 197 | """The input of normlization will be (M, C, K), where M is batch size, 198 | C is channel size and K is sequence length. 199 | """ 200 | if norm_type == "gLN": 201 | return GlobalLayerNorm(channel_size) 202 | elif norm_type == "cLN": 203 | return ChannelwiseLayerNorm(channel_size) 204 | else: # norm_type == "BN": 205 | # Given input (M, C, K), nn.BatchNorm1d(C) will accumulate statics 206 | # along M and K, so this BN usage is right. 207 | return nn.BatchNorm1d(channel_size) 208 | 209 | 210 | # TODO: Use nn.LayerNorm to impl cLN to speed up 211 | class ChannelwiseLayerNorm(nn.Module): 212 | """Channel-wise Layer Normalization (cLN)""" 213 | def __init__(self, channel_size): 214 | super(ChannelwiseLayerNorm, self).__init__() 215 | self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1)) # [1, N, 1] 216 | self.beta = nn.Parameter(torch.Tensor(1, channel_size,1 )) # [1, N, 1] 217 | self.reset_parameters() 218 | 219 | def reset_parameters(self): 220 | self.gamma.data.fill_(1) 221 | self.beta.data.zero_() 222 | 223 | def forward(self, y): 224 | """ 225 | Args: 226 | y: [M, N, K], M is batch size, N is channel size, K is length 227 | Returns: 228 | cLN_y: [M, N, K] 229 | """ 230 | mean = torch.mean(y, dim=1, keepdim=True) # [M, 1, K] 231 | var = torch.var(y, dim=1, keepdim=True, unbiased=False) # [M, 1, K] 232 | cLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta 233 | return cLN_y 234 | 235 | 236 | class GlobalLayerNorm(nn.Module): 237 | """Global Layer Normalization (gLN)""" 238 | def __init__(self, channel_size): 239 | super(GlobalLayerNorm, self).__init__() 240 | self.gamma = nn.Parameter(torch.Tensor(1, channel_size, 1)) # [1, N, 1] 241 | self.beta = nn.Parameter(torch.Tensor(1, channel_size,1 )) # [1, N, 1] 242 | self.reset_parameters() 243 | 244 | def reset_parameters(self): 245 | self.gamma.data.fill_(1) 246 | self.beta.data.zero_() 247 | 248 | def forward(self, y): 249 | """ 250 | Args: 251 | y: [M, N, K], M is batch size, N is channel size, K is length 252 | Returns: 253 | gLN_y: [M, N, K] 254 | """ 255 | # TODO: in torch 1.0, torch.mean() support dim list 256 | mean = y.mean(dim=1, keepdim=True).mean(dim=2, keepdim=True) #[M, 1, 1] 257 | var = (torch.pow(y-mean, 2)).mean(dim=1, keepdim=True).mean(dim=2, keepdim=True) 258 | gLN_y = self.gamma * (y - mean) / torch.pow(var + EPS, 0.5) + self.beta 259 | return gLN_y 260 | 261 | 262 | --------------------------------------------------------------------------------