├── Modelsave ├── decoder.pth.tar └── encoder.pth.tar ├── README.md ├── Model_evaluation_encoder.py ├── Model_evaluation_decoder.py ├── Model_train.py └── Model_define_pytorch.py /Modelsave/decoder.pth.tar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xufana7/AutoEncoder-with-pytorch/HEAD/Modelsave/decoder.pth.tar -------------------------------------------------------------------------------- /Modelsave/encoder.pth.tar: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/xufana7/AutoEncoder-with-pytorch/HEAD/Modelsave/encoder.pth.tar -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # This is the code for competition AI+无线通信 2020 NIAC https://naic.pcl.ac.cn/ with pytorch. 2 | 3 | ## requirements 4 | - pytorch > 0.4 5 | - numpy 6 | - h5py 7 | ## result 8 | test score = 0.67 9 | -------------------------------------------------------------------------------- /Model_evaluation_encoder.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import numpy as np 3 | import h5py 4 | from Model_define_pytorch import AutoEncoder, DatasetFolder 5 | import torch 6 | import os 7 | 8 | 9 | # Parameters for training 10 | os.environ["CUDA_VISIBLE_DEVICES"] = '0' 11 | batch_size = 64 12 | num_workers = 4 13 | # Data parameters 14 | img_height = 16 15 | img_width = 32 16 | img_channels = 2 17 | feedback_bits = 128 18 | model_ID = 'CsiNet' # Model Number 19 | 20 | # load test data 21 | mat = h5py.File('./data/H_test.mat') 22 | data = np.transpose(mat['H_test'] ) 23 | data = data.astype('float32') 24 | data = np.reshape(data, [len(data), img_channels, img_height, img_width]) 25 | 26 | # load model 27 | model = AutoEncoder(feedback_bits).cuda() 28 | model_encoder = model.encoder 29 | model_path = './Modelsave/encoder.pth.tar' 30 | model_encoder.load_state_dict(torch.load(model_path)['state_dict']) 31 | print("weight loaded") 32 | 33 | #dataLoader for test 34 | test_dataset = DatasetFolder(data) 35 | test_loader = torch.utils.data.DataLoader( 36 | test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) 37 | 38 | # test 39 | model_encoder.eval() 40 | encode_feature = [] 41 | with torch.no_grad(): 42 | for i, input in enumerate(test_loader): 43 | # convert numpy to Tensor 44 | input = input.cuda() 45 | output = model_encoder(input) 46 | output = output.cpu().numpy() 47 | if i == 0: 48 | encode_feature = output 49 | else: 50 | encode_feature = np.concatenate((encode_feature, output), axis=0) 51 | print("feedbackbits length is ", np.shape(encode_feature)[-1]) 52 | np.save('./Modelsave/encoder_output.npy', encode_feature) 53 | -------------------------------------------------------------------------------- /Model_evaluation_decoder.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | import numpy as np 3 | import h5py 4 | from Model_define_pytorch import NMSE, AutoEncoder, DatasetFolder 5 | import torch 6 | import os 7 | 8 | # Parameters for training 9 | os.environ["CUDA_VISIBLE_DEVICES"] = '0' 10 | batch_size = 64 11 | num_workers = 4 12 | # parameter setting 13 | img_height = 16 14 | img_width = 32 15 | img_channels = 2 16 | feedback_bits = 128 17 | # Data loading 18 | data_load_address = './data' 19 | mat = h5py.File(data_load_address+'/H_test.mat', 'r') 20 | data = np.transpose(mat['H_test']) 21 | data = data.astype('float32') 22 | x_test = np.reshape(data, [len(data), img_channels, img_height, img_width]) 23 | 24 | # load encoder_output 25 | decode_input = np.load('./Modelsave/encoder_output.npy') 26 | 27 | # load model and test NMSE 28 | model = AutoEncoder(feedback_bits).cuda() 29 | model_decoder = model.decoder 30 | model_path = './Modelsave/decoder.pth.tar' 31 | model_decoder.load_state_dict(torch.load(model_path)['state_dict']) 32 | print("weight loaded") 33 | 34 | # dataLoader for test 35 | test_dataset = DatasetFolder(decode_input) 36 | test_loader = torch.utils.data.DataLoader( 37 | test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) 38 | 39 | # test 40 | model_decoder.eval() 41 | y_test = [] 42 | with torch.no_grad(): 43 | for i, input in enumerate(test_loader): 44 | # convert numpy to Tensor 45 | input = input.cuda() 46 | output = model_decoder(input) 47 | output = output.cpu().numpy() 48 | if i == 0: 49 | y_test = output 50 | else: 51 | y_test = np.concatenate((y_test, output), axis=0) 52 | 53 | # need convert channel first to channel last for evaluate. 54 | print('The NMSE is ' + np.str(NMSE(np.transpose(x_test, (0, 2, 3, 1)), np.transpose(y_test, (0, 2, 3, 1))))) 55 | -------------------------------------------------------------------------------- /Model_train.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """An Implement of an autoencoder with pytorch. 3 | This is the template code for 2020 NIAC https://naic.pcl.ac.cn/. 4 | The code is based on the sample code with tensorflow for 2020 NIAC and it can only run with GPUS. 5 | If you have any questions, please contact me with https://github.com/xufana7/AutoEncoder-with-pytorch 6 | Author, Fan xu Aug 2020 7 | """ 8 | import numpy as np 9 | import h5py 10 | import torch 11 | from Model_define_pytorch import AutoEncoder, DatasetFolder 12 | import os 13 | import torch.nn as nn 14 | 15 | # Parameters for training 16 | os.environ["CUDA_VISIBLE_DEVICES"] = '0' 17 | use_single_gpu = True # select whether using single gpu or multiple gpus 18 | torch.manual_seed(1) 19 | batch_size = 512 20 | epochs = 1000 21 | learning_rate = 1e-3 22 | num_workers = 4 23 | print_freq = 100 # print frequency (default: 60) 24 | # parameters for data 25 | feedback_bits = 128 26 | img_height = 16 27 | img_width = 32 28 | img_channels = 2 29 | 30 | # Model construction 31 | model = AutoEncoder(feedback_bits) 32 | if use_single_gpu: 33 | model = model.cuda() 34 | 35 | else: 36 | # DataParallel will divide and allocate batch_size to all available GPUs 37 | autoencoder = torch.nn.DataParallel(model).cuda() 38 | 39 | criterion = nn.MSELoss().cuda() 40 | optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) 41 | 42 | # Data loading 43 | data_load_address = './data' 44 | mat = h5py.File(data_load_address + '/H_train.mat', 'r') 45 | data = np.transpose(mat['H_train']) # shape=(320000, 1024) 46 | data = data.astype('float32') 47 | data = np.reshape(data, [len(data), img_channels, img_height, img_width]) 48 | np.random.shuffle(data) 49 | start = int(data.shape[0] * 0.7) 50 | x_train, x_test = data[:start], data[start:] 51 | 52 | # dataLoader for training 53 | train_dataset = DatasetFolder(x_train) 54 | train_loader = torch.utils.data.DataLoader( 55 | train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers, pin_memory=True) 56 | 57 | # dataLoader for training 58 | test_dataset = DatasetFolder(x_test) 59 | test_loader = torch.utils.data.DataLoader( 60 | test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, pin_memory=True) 61 | best_loss = 1 62 | for epoch in range(epochs): 63 | # model training 64 | model.train() 65 | for i, input in enumerate(train_loader): 66 | # adjust learning rate 67 | if epoch == 300: 68 | for param_group in optimizer.param_groups: 69 | param_group['lr'] = learning_rate * 0.1 70 | input = input.cuda() 71 | # compute output 72 | output = model(input) 73 | loss = criterion(output, input) 74 | # compute gradient and do Adam step 75 | optimizer.zero_grad() 76 | loss.backward() 77 | optimizer.step() 78 | if i % print_freq == 0: 79 | print('Epoch: [{0}][{1}/{2}]\t' 80 | 'Loss {loss:.4f}\t'.format( 81 | epoch, i, len(train_loader), loss=loss.item())) 82 | # model evaluating 83 | model.eval() 84 | total_loss = 0 85 | with torch.no_grad(): 86 | for i, input in enumerate(test_loader): 87 | input = input.cuda() 88 | output = model(input) 89 | total_loss += criterion(output, input).item() * input.size(0) 90 | average_loss = total_loss / len(test_dataset) 91 | if average_loss < best_loss: 92 | # model save 93 | # save encoder 94 | modelSave1 = './Modelsave/encoder.pth.tar' 95 | torch.save({'state_dict': model.encoder.state_dict(), }, modelSave1) 96 | # save decoder 97 | modelSave2 = './Modelsave/decoder.pth.tar' 98 | torch.save({'state_dict': model.decoder.state_dict(), }, modelSave2) 99 | print("Model saved") 100 | best_loss = average_loss 101 | -------------------------------------------------------------------------------- /Model_define_pytorch.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | """An Implement of an autoencoder with pytorch. 3 | This is the template code for 2020 NIAC https://naic.pcl.ac.cn/. 4 | The code is based on the sample code with tensorflow for 2020 NIAC and it can only run with GPUS. 5 | If you have any questions, please contact me with https://github.com/xufana7/AutoEncoder-with-pytorch 6 | Author, Fan xu Aug 2020 7 | """ 8 | import numpy as np 9 | import torch.nn as nn 10 | import torch 11 | import torch.nn.functional as F 12 | from torch.utils.data import Dataset 13 | 14 | 15 | # This part implement the quantization and dequantization operations. 16 | # The output of the encoder must be the bitstream. 17 | def Num2Bit(Num, B): 18 | Num_ = Num.type(torch.uint8) 19 | 20 | def integer2bit(integer, num_bits=B * 2): 21 | dtype = integer.type() 22 | exponent_bits = -torch.arange(-(num_bits - 1), 1).type(dtype) 23 | exponent_bits = exponent_bits.repeat(integer.shape + (1,)) 24 | out = integer.unsqueeze(-1) // 2 ** exponent_bits 25 | return (out - (out % 1)) % 2 26 | 27 | bit = integer2bit(Num_) 28 | bit = (bit[:, :, B:]).reshape(-1, Num_.shape[1] * B) 29 | return bit.type(torch.float32) 30 | 31 | 32 | def Bit2Num(Bit, B): 33 | Bit_ = Bit.type(torch.float32) 34 | Bit_ = torch.reshape(Bit_, [-1, int(Bit_.shape[1] / B), B]) 35 | num = torch.zeros(Bit_[:, :, 1].shape).cuda() 36 | for i in range(B): 37 | num = num + Bit_[:, :, i] * 2 ** (B - 1 - i) 38 | return num 39 | 40 | 41 | class Quantization(torch.autograd.Function): 42 | @staticmethod 43 | def forward(ctx, x, B): 44 | ctx.constant = B 45 | step = 2 ** B 46 | out = torch.round(x * step - 0.5) 47 | out = Num2Bit(out, B) 48 | return out 49 | 50 | @staticmethod 51 | def backward(ctx, grad_output): 52 | # return as many input gradients as there were arguments. 53 | # Gradients of constant arguments to forward must be None. 54 | # Gradient of a number is the sum of its four bits. 55 | b, _ = grad_output.shape 56 | grad_num = torch.sum(grad_output.reshape(b, -1, ctx.constant), dim=2) 57 | return grad_num, None 58 | 59 | 60 | class Dequantization(torch.autograd.Function): 61 | @staticmethod 62 | def forward(ctx, x, B): 63 | ctx.constant = B 64 | step = 2 ** B 65 | out = Bit2Num(x, B) 66 | out = (out + 0.5) / step 67 | return out 68 | 69 | @staticmethod 70 | def backward(ctx, grad_output): 71 | # return as many input gradients as there were arguments. 72 | # Gradients of non-Tensor arguments to forward must be None. 73 | # repeat the gradient of a Num for four time. 74 | b, c = grad_output.shape 75 | grad_output = grad_output.unsqueeze(2) / ctx.constant 76 | grad_bit = grad_output.expand(b, c, ctx.constant) 77 | return torch.reshape(grad_bit, (-1, c * ctx.constant)), None 78 | 79 | 80 | class QuantizationLayer(nn.Module): 81 | 82 | def __init__(self, B): 83 | super(QuantizationLayer, self).__init__() 84 | self.B = B 85 | 86 | def forward(self, x): 87 | out = Quantization.apply(x, self.B) 88 | return out 89 | 90 | 91 | class DequantizationLayer(nn.Module): 92 | 93 | def __init__(self, B): 94 | super(DequantizationLayer, self).__init__() 95 | self.B = B 96 | 97 | def forward(self, x): 98 | out = Dequantization.apply(x, self.B) 99 | return out 100 | 101 | 102 | def conv3x3(in_planes, out_planes, stride=1): 103 | """3x3 convolution with padding""" 104 | return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, 105 | padding=1, bias=True) 106 | 107 | 108 | class Encoder(nn.Module): 109 | B = 4 110 | 111 | def __init__(self, feedback_bits): 112 | super(Encoder, self).__init__() 113 | self.conv1 = conv3x3(2, 2) 114 | self.conv2 = conv3x3(2, 2) 115 | self.fc = nn.Linear(1024, int(feedback_bits / self.B)) 116 | self.sig = nn.Sigmoid() 117 | self.quantize = QuantizationLayer(self.B) 118 | 119 | def forward(self, x): 120 | out = F.relu(self.conv1(x)) 121 | out = F.relu(self.conv2(out)) 122 | out = out.view(-1, 1024) 123 | out = self.fc(out) 124 | out = self.sig(out) 125 | out = self.quantize(out) 126 | 127 | return out 128 | 129 | 130 | class Decoder(nn.Module): 131 | B = 4 132 | 133 | def __init__(self, feedback_bits): 134 | super(Decoder, self).__init__() 135 | self.feedback_bits = feedback_bits 136 | self.dequantize = DequantizationLayer(self.B) 137 | self.multiConvs = nn.ModuleList() 138 | self.fc = nn.Linear(int(feedback_bits / self.B), 1024) 139 | self.out_cov = conv3x3(2, 2) 140 | self.sig = nn.Sigmoid() 141 | 142 | for _ in range(3): 143 | self.multiConvs.append(nn.Sequential( 144 | conv3x3(2, 8), 145 | nn.ReLU(), 146 | conv3x3(8, 16), 147 | nn.ReLU(), 148 | conv3x3(16, 2), 149 | nn.ReLU())) 150 | 151 | def forward(self, x): 152 | out = self.dequantize(x) 153 | out = out.view(-1, int(self.feedback_bits / self.B)) 154 | out = self.sig(self.fc(out)) 155 | out = out.view(-1, 2, 16, 32) 156 | for i in range(3): 157 | residual = out 158 | out = self.multiConvs[i](out) 159 | out = residual + out 160 | 161 | out = self.out_cov(out) 162 | out = self.sig(out) 163 | return out 164 | 165 | 166 | # Note: Do not modify following class and keep it in your submission. 167 | # feedback_bits is 128 by default. 168 | class AutoEncoder(nn.Module): 169 | 170 | def __init__(self, feedback_bits): 171 | super(AutoEncoder, self).__init__() 172 | self.encoder = Encoder(feedback_bits) 173 | self.decoder = Decoder(feedback_bits) 174 | 175 | def forward(self, x): 176 | feature = self.encoder(x) 177 | out = self.decoder(feature) 178 | return out 179 | 180 | 181 | def NMSE(x, x_hat): 182 | x_real = np.reshape(x[:, :, :, 0], (len(x), -1)) 183 | x_imag = np.reshape(x[:, :, :, 1], (len(x), -1)) 184 | x_hat_real = np.reshape(x_hat[:, :, :, 0], (len(x_hat), -1)) 185 | x_hat_imag = np.reshape(x_hat[:, :, :, 1], (len(x_hat), -1)) 186 | x_C = x_real - 0.5 + 1j * (x_imag - 0.5) 187 | x_hat_C = x_hat_real - 0.5 + 1j * (x_hat_imag - 0.5) 188 | power = np.sum(abs(x_C) ** 2, axis=1) 189 | mse = np.sum(abs(x_C - x_hat_C) ** 2, axis=1) 190 | nmse = np.mean(mse / power) 191 | return nmse 192 | 193 | 194 | def Score(NMSE): 195 | score = 1 - NMSE 196 | return score 197 | 198 | 199 | # dataLoader 200 | class DatasetFolder(Dataset): 201 | 202 | def __init__(self, matData): 203 | self.matdata = matData 204 | 205 | def __getitem__(self, index): 206 | return self.matdata[index] 207 | 208 | def __len__(self): 209 | return self.matdata.shape[0] 210 | --------------------------------------------------------------------------------