import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import matplotlib.pyplot as plt
# local model
import sys
sys.path.append("../network")
import Coral
from lstm import LSTMHardSigmoid
from AdaBN import AdaBN
sys.path.append("../network/AutoEncoder")
import AutoEncoder

class cnnblstm_with_adabn(nn.Module):
	PARAMS_FILE = "params.pkl"
	PARAMS_AE = "params_ae.pkl"
	NET1_ADABN = "net1_adabn"
	NET2_ADABN = "net2_adabn"
	NET3_ADABN = "net3_adabn"

	def __init__(self, time_steps = 800, n_features = 3, n_outputs = 10, use_cuda = False, params_dir = "./params", enable_CORAL = False):
		super(cnnblstm_with_adabn, self).__init__()

		self.time_steps = time_steps
		self.n_features = n_features
		self.n_outputs = n_outputs
		self.use_cuda = use_cuda
		self.params_dir = params_dir
		if not os.path.exists(self.params_dir):
			os.mkdir(self.params_dir)
		self.enable_CORAL = enable_CORAL

		self.n_filters = 128
		self.kernel_size = 15
		self.n_hidden = 150	# 150
		self.n_layers = 1
		self.bidirectional = True

		# self.ae = AutoEncoder.load_AE(type = "ConvAE", time_steps = self.time_steps, n_features = self.n_features, use_cuda = self.use_cuda, params_pkl = os.path.join(self.params_dir, cnnblstm_with_adabn.PARAMS_AE))

		# build net1 cnn
		self.net1 = nn.Sequential(
			nn.Conv1d(in_channels = self.n_features, out_channels = self.n_filters, kernel_size = self.kernel_size),
			# nn.Conv1d(in_channels = self.ae.n_filters3, out_channels = self.n_filters, kernel_size = self.kernel_size),
			nn.ReLU(),
			# nn.Sigmoid(),
			nn.Dropout(p = 0.5),
			nn.MaxPool1d(kernel_size = 2)
		)

		# build net1_adabn
		self.net1_adabn = AdaBN(self.n_filters, variables_dir = os.path.join(self.params_dir, cnnblstm_with_adabn.NET1_ADABN), use_cuda = self.use_cuda)

		# build net2 blstm
		# self.net2 = nn.LSTM(input_size = self.n_filters, hidden_size = self.n_hidden, num_layers = self.n_layers, dropout = 0.2, batch_first = True, bidirectional = self.bidirectional, bias = True)
		self.net2 = LSTMHardSigmoid(input_size = self.n_filters, hidden_size = self.n_hidden, num_layers = self.n_layers, dropout = 0.2, batch_first = True, bidirectional = self.bidirectional, bias = True)

		# build net2_adabn
		if self.bidirectional:
			n_blstm_output = self.n_hidden * 2
		else:
			n_blstm_output = self.n_hidden
		self.net2_adabn = AdaBN(n_blstm_output, variables_dir = os.path.join(self.params_dir, cnnblstm_with_adabn.NET2_ADABN), use_cuda = self.use_cuda)

		# build net3 fc
		self.net3 = nn.Sequential(
			nn.Linear(n_blstm_output, 50, bias = True),
			nn.ReLU(),
			# nn.Sigmoid(),
		)

		# build net3_adabn
		self.net3_adabn = AdaBN(50, variables_dir = os.path.join(self.params_dir, cnnblstm_with_adabn.NET3_ADABN), use_cuda = self.use_cuda)

		# build net4 fc
		self.net4 = nn.Sequential(
			nn.Dropout(p = 0.2),
			nn.Linear(50, self.n_outputs, bias = True),
			nn.Softmax(dim = 1)
		)

	def init_hidden(self, batch_size):
		"""
		init blstm's hidden states
		"""
		if self.bidirectional:
			n_layers = self.n_layers * 2
		else:
			n_layers = self.n_layers
		if self.use_cuda:
			hidden_state = torch.zeros(n_layers, batch_size, self.n_hidden).cuda()
			cell_state = torch.zeros(n_layers, batch_size, self.n_hidden).cuda()
		else:
			hidden_state = torch.zeros(n_layers, batch_size, self.n_hidden)
			cell_state = torch.zeros(n_layers, batch_size, self.n_hidden)
		self.hidden = (hidden_state, cell_state)

	def reset_parameters(self):
		"""
		temp useless
		Here we reproduce Keras default initialization weights for consistency with Keras version
		"""
		# get weights & bias set
		net1_weights = ((name, param.data) for name, param in self.named_parameters() if (("weight" in name) and (("net1" in name) and ("net1_adabn" not in name))))
		net1_biases = ((name, param.data) for name, param in self.named_parameters() if (("bias" in name) and (("net1" in name) and ("net1_adabn" not in name))))
		# net2_weights = ((name, param.data) for name, param in self.named_parameters() if (("weight" in name) and (("net2" in name) and ("net2_adabn" not in name))))
		# net2_biases = ((name, param.data) for name, param in self.named_parameters() if (("bias" in name) and (("net2" in name) and ("net2_adabn" not in name))))
		net3_weights = ((name, param.data) for name, param in self.named_parameters() if (("weight" in name) and (("net3" in name) and ("net3_adabn" not in name))))
		net3_biases = ((name, param.data) for name, param in self.named_parameters() if (("bias" in name) and (("net3" in name) and ("net3_adabn" not in name))))
		net4_weights = ((name, param.data) for name, param in self.named_parameters() if (("weight" in name) and (("net4" in name) and ("net4_adabn" not in name))))
		net4_biases = ((name, param.data) for name, param in self.named_parameters() if (("bias" in name) and (("net4" in name) and ("net4_adabn" not in name))))
		# init weights & bias
		# self.ae.reset_parameters()
		for name, params_data in net1_weights:
			# print(name)
			nn.init.xavier_uniform_(params_data)
		for name, params_data in net1_biases:
			nn.init.constant_(params_data, 0)
		self.net1_adabn.reset_parameters()
		self.net2.reset_parameters()		# lstm reset parameters
		self.net2_adabn.reset_parameters()
		for name, params_data in net3_weights:
			nn.init.xavier_uniform_(params_data)
		for name, params_data in net3_biases:
			nn.init.constant_(params_data, 0)
		self.net3_adabn.reset_parameters()
		for name, params_data in net4_weights:
			nn.init.xavier_uniform_(params_data)
		for name, params_data in net4_biases:
			nn.init.constant_(params_data, 0)

	def forward(self, input):
		"""
		compute the output of input according to the entire network model
		"""
		# print(input.shape)
		# AutoEncoder
		# input = self.ae.encoder(input)
		# input = self.ae(input)
		# MaxPool1d
		maxPool1d_output = self.net1(input)
		# maxPool1d_adabn_output = maxPool1d_output
		maxPool1d_adabn_output, maxPool1d_output = self.net1_adabn(maxPool1d_output), None
		maxPool1d_adabn_t_output = maxPool1d_adabn_output.permute(0, 2, 1).contiguous()
		# BiLSTM
		(bilstm_output, _), maxPool1d_adabn_t_output = self.net2(maxPool1d_adabn_t_output, None), None
		# MaxPooling1D time_steps
		bilstm_output = bilstm_output.permute(0, 2, 1)
		maxPooling_output, bilstm_output = F.max_pool1d(bilstm_output, kernel_size = bilstm_output.size(2)).squeeze(2), None
		# maxPooling_adabn_output = maxPooling_output
		maxPooling_adabn_output, maxPooling_output = self.net2_adabn(maxPooling_output), None
		# get classifier
		net3_output, maxPooling_adabn_output = self.net3(maxPooling_adabn_output), None
		net3_adabn_output, net3_output = self.net3_adabn(net3_output), None
		linear2_softmax_output, net3_adabn_output = self.net4(net3_adabn_output), None

		return linear2_softmax_output

	def update_adabn_running_stats(self):
		"""
		update adabn running states, update mu_j with mu_j_next to start next round
		"""
		self.net1_adabn.update_running_stats()
		self.net2_adabn.update_running_stats()
		self.net3_adabn.update_running_stats()

	def trainAllLayers(self, train_x, train_y, test_x = None, learning_rate = 0.001, n_epoches = 20, batch_size = 20, shuffle = True):
		"""
		train all layers of network model
		"""
		# print(os.environ["CUDA_VISIBLE_DEVICES"])
		# CORAL
		if self.enable_CORAL:
			if test_x == None:
				print("ERROR: (in cnnblstm_with_adabn.trainAllLayers) test_x == None!")
				return
			# review train_x & test_x
			train_x = train_x.view(-1, self.time_steps * self.n_features)
			test_x = test_x.view(-1, self.time_steps * self.n_features)
			# get CORAL(train_x, test_x)
			train_x = Coral.CORAL_torch(train_x, test_x)
		# review train_x
		train_x = train_x.view(-1, self.n_features, self.time_steps)
		# optimize all cnn parameters
		params = [{"params": model.parameters()} for model in self.children() if model  not in [self.ae]]
		optimizer = torch.optim.Adam(params, lr = learning_rate)
		# the target label is not one-hotted
		loss_func = nn.CrossEntropyLoss()

		# init params
		self.reset_parameters()

		# load params
		self.load_params()

		# set train mode True
		self.train()

		# get parallel model
		parallel_cba = self
		if self.use_cuda:
			# print("we use cuda!")
			parallel_cba = torch.nn.DataParallel(self, device_ids = range(torch.cuda.device_count()))
			# parallel_cba = parallel_cba.cuda()

		# if use_cuda
		if self.use_cuda:
			train_x = train_x.cuda()
			train_y = train_y.cuda()

		"""
		# get autoencoder
		self.ae = AutoEncoder.train_AE(self.ae, train_x, train_x, n_epoches = 20)
		self.ae.save_params()
		"""
 
		# get train_data
		train_data = torch.utils.data.TensorDataset(train_x, train_y)
		# Data Loader for easy mini-batch return in training
		train_loader = torch.utils.data.DataLoader(dataset = train_data, batch_size = batch_size, shuffle = shuffle)

		# training and testing
		for epoch in range(n_epoches):
			# init loss & acc
			train_loss = 0
			train_acc = 0
			for step, (b_x, b_y) in enumerate(train_loader):		# gives batch data
				b_x = b_x.view(-1, self.n_features, self.time_steps)	# reshape x to (batch, n_features, time_step)
				if self.use_cuda:
					b_x, b_y = Variable(b_x).cuda(), Variable(b_y).cuda()
				else:
					b_x, b_y = Variable(b_x), Variable(b_y)
				"""
				# get hidden
				if self.use_cuda:
					self.init_hidden(b_x.size(0) // torch.cuda.device_count())
				else:
					self.init_hidden(b_x.size(0))
				"""
				# update adabn running stats
				self.update_adabn_running_stats()
				# get output
				output = parallel_cba(b_x)									# CNN_BLSTM output
				# get loss
				loss = loss_func(output, b_y)						# cross entropy loss
				train_loss += loss.item() * len(b_y)
				_, pre = torch.max(output, 1)
				num_acc = (pre == b_y).sum()
				train_acc += num_acc.item()
				# backward
				optimizer.zero_grad()								# clear gradients for this training step
				loss.backward()										# backpropagation, compute gradients
				optimizer.step()									# apply gradients

				# print loss
				# if (step + 1) % 5 == 0:
					# print("[{}/{}], train loss is: {:.6f}, train acc is: {:.6f}".format(step, len(train_loader), train_loss / ((step + 1) * batch_size), train_acc / ((step + 1) * batch_size)))
			print("[{}/{}], train loss is: {:.6f}, train acc is: {:.6f}".format(len(train_loader), len(train_loader), train_loss / (len(train_loader) * batch_size), train_acc / (len(train_loader) * batch_size)))

			# save params
			self.save_params()

		# print("train finish!")

	def getTestAccuracy(self, test_x, test_y):
		"""
		test network model with test set
		"""
		# init params
		self.reset_parameters()

		# load params
		self.load_params()

		# set eval
		self.eval()

		# get parallel model
		parallel_cba = self
		if self.use_cuda:
			# print("we use cuda!")
			parallel_cba = torch.nn.DataParallel(self, device_ids = range(torch.cuda.device_count()))
			# parallel_cba = parallel_cba.cuda()
		# cuda test_data
		with torch.no_grad():
			if self.use_cuda:
				test_x, test_y = Variable(test_x).cuda(), Variable(test_y).cuda()
			else:
				test_x, test_y = Variable(test_x), Variable(test_y)
		"""
		# get hidden
		if self.use_cuda:
			self.init_hidden(test_x.size(0) // torch.cuda.device_count())
		else:
			self.init_hidden(test_x.size(0))
		"""
		# update adabn running stats
		self.update_adabn_running_stats()
		# get output
		with torch.no_grad():
			output = parallel_cba(test_x)
		# print(output)
		prediction = torch.max(output, 1)[1]
		pred_y = prediction.cpu().data.numpy()
		# print(pred_y)
		target_y = test_y.cpu().data.numpy()
		# print(test_y)

		accuracy = float((pred_y == target_y).astype(int).sum()) / float(target_y.size)
		# print("Accuracy: ", str(accuracy))
		return accuracy

	def save_params(self):
		"""
		save params & adabn's inner stats
		"""
		self.save_adabn_variables()
		torch.save(self.state_dict(), os.path.join(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE))
		# self.ae.save_params()
		# print("save_params success!")

	def save_adabn_variables(self):
		"""
		save adabn's inner stats
		"""
		self.net1_adabn.save_attrs()
		self.net2_adabn.save_attrs()
		self.net3_adabn.save_attrs()

	def load_params(self):
		"""
		load params & adabn's inner stats
		"""
		self.load_adabn_variables()
		if os.path.exists(os.path.join(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE)):
			if self.use_cuda:
				self.load_state_dict(torch.load(os.path.join(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE), map_location = torch.device('cuda')))
			else:
				self.load_state_dict(torch.load(os.path.join(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE), map_location = torch.device('cpu')))
			# print("load_params success!")
		# self.ae.load_params()

	def load_adabn_variables(self):
		"""
		load adabn's inner stats
		"""
		self.net1_adabn.load_attrs()
		self.net2_adabn.load_attrs()
		self.net3_adabn.load_attrs()

	def get_model(self, pre_trained = False):
		"""
		get pretrained model
		"""
		if pre_trained:
			self.load_params()
		return self

if __name__ == '__main__':
	use_cuda = torch.cuda.is_available()
	if use_cuda:
		cnnblstm = cnnblstm_with_adabn(use_cuda = use_cuda).cuda()
	else:
		cnnblstm = cnnblstm_with_adabn(use_cuda = use_cuda)
	print(cnnblstm)
	# get train_x, train_y
	train_x = torch.rand(20, 3, 800, dtype = torch.float32)
	train_y = torch.randint(10, (20, ), dtype = torch.int64)
	# train_y = torch.LongTensor(20, 1).random_() % 10
	print(train_x.type())
	# train_y = torch.zeros(20, 10).scatter_(1, train_y, 1)
	print(train_y)
	train_data = torch.utils.data.TensorDataset(train_x, train_y)
	cnnblstm.trainAllLayers(train_data)

[((189, 218), 'sys.path.append', 'sys.path.append', (['"""../network"""'], {}), "('../network')\n", (204, 218), False, 'import sys\n'), ((289, 330), 'sys.path.append', 'sys.path.append', (['"""../network/AutoEncoder"""'], {}), "('../network/AutoEncoder')\n", (304, 330), False, 'import sys\n'), ((13068, 13093), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (13091, 13093), False, 'import torch\n'), ((13282, 13325), 'torch.rand', 'torch.rand', (['(20)', '(3)', '(800)'], {'dtype': 'torch.float32'}), '(20, 3, 800, dtype=torch.float32)\n', (13292, 13325), False, 'import torch\n'), ((13339, 13382), 'torch.randint', 'torch.randint', (['(10)', '(20,)'], {'dtype': 'torch.int64'}), '(10, (20,), dtype=torch.int64)\n', (13352, 13382), False, 'import torch\n'), ((13548, 13596), 'torch.utils.data.TensorDataset', 'torch.utils.data.TensorDataset', (['train_x', 'train_y'], {}), '(train_x, train_y)\n', (13578, 13596), False, 'import torch\n'), ((2066, 2246), 'lstm.LSTMHardSigmoid', 'LSTMHardSigmoid', ([], {'input_size': 'self.n_filters', 'hidden_size': 'self.n_hidden', 'num_layers': 'self.n_layers', 'dropout': '(0.2)', 'batch_first': '(True)', 'bidirectional': 'self.bidirectional', 'bias': '(True)'}), '(input_size=self.n_filters, hidden_size=self.n_hidden,\n    num_layers=self.n_layers, dropout=0.2, batch_first=True, bidirectional=\n    self.bidirectional, bias=True)\n', (2081, 2246), False, 'from lstm import LSTMHardSigmoid\n'), ((7921, 7963), 'torch.optim.Adam', 'torch.optim.Adam', (['params'], {'lr': 'learning_rate'}), '(params, lr=learning_rate)\n', (7937, 7963), False, 'import torch\n'), ((8019, 8040), 'torch.nn.CrossEntropyLoss', 'nn.CrossEntropyLoss', ([], {}), '()\n', (8038, 8040), True, 'import torch.nn as nn\n'), ((8650, 8698), 'torch.utils.data.TensorDataset', 'torch.utils.data.TensorDataset', (['train_x', 'train_y'], {}), '(train_x, train_y)\n', (8680, 8698), False, 'import torch\n'), ((8771, 8862), 'torch.utils.data.DataLoader', 'torch.utils.data.DataLoader', ([], {'dataset': 'train_data', 'batch_size': 'batch_size', 'shuffle': 'shuffle'}), '(dataset=train_data, batch_size=batch_size,\n    shuffle=shuffle)\n', (8798, 8862), False, 'import torch\n'), ((869, 900), 'os.path.exists', 'os.path.exists', (['self.params_dir'], {}), '(self.params_dir)\n', (883, 900), False, 'import os\n'), ((905, 930), 'os.mkdir', 'os.mkdir', (['self.params_dir'], {}), '(self.params_dir)\n', (913, 930), False, 'import os\n'), ((1355, 1456), 'torch.nn.Conv1d', 'nn.Conv1d', ([], {'in_channels': 'self.n_features', 'out_channels': 'self.n_filters', 'kernel_size': 'self.kernel_size'}), '(in_channels=self.n_features, out_channels=self.n_filters,\n    kernel_size=self.kernel_size)\n', (1364, 1456), True, 'import torch.nn as nn\n'), ((1576, 1585), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (1583, 1585), True, 'import torch.nn as nn\n'), ((1609, 1626), 'torch.nn.Dropout', 'nn.Dropout', ([], {'p': '(0.5)'}), '(p=0.5)\n', (1619, 1626), True, 'import torch.nn as nn\n'), ((1633, 1660), 'torch.nn.MaxPool1d', 'nn.MaxPool1d', ([], {'kernel_size': '(2)'}), '(kernel_size=2)\n', (1645, 1660), True, 'import torch.nn as nn\n'), ((2577, 2617), 'torch.nn.Linear', 'nn.Linear', (['n_blstm_output', '(50)'], {'bias': '(True)'}), '(n_blstm_output, 50, bias=True)\n', (2586, 2617), True, 'import torch.nn as nn\n'), ((2624, 2633), 'torch.nn.ReLU', 'nn.ReLU', ([], {}), '()\n', (2631, 2633), True, 'import torch.nn as nn\n'), ((2866, 2883), 'torch.nn.Dropout', 'nn.Dropout', ([], {'p': '(0.2)'}), '(p=0.2)\n', (2876, 2883), True, 'import torch.nn as nn\n'), ((2890, 2930), 'torch.nn.Linear', 'nn.Linear', (['(50)', 'self.n_outputs'], {'bias': '(True)'}), '(50, self.n_outputs, bias=True)\n', (2899, 2930), True, 'import torch.nn as nn\n'), ((2937, 2954), 'torch.nn.Softmax', 'nn.Softmax', ([], {'dim': '(1)'}), '(dim=1)\n', (2947, 2954), True, 'import torch.nn as nn\n'), ((3324, 3372), 'torch.zeros', 'torch.zeros', (['n_layers', 'batch_size', 'self.n_hidden'], {}), '(n_layers, batch_size, self.n_hidden)\n', (3335, 3372), False, 'import torch\n'), ((3389, 3437), 'torch.zeros', 'torch.zeros', (['n_layers', 'batch_size', 'self.n_hidden'], {}), '(n_layers, batch_size, self.n_hidden)\n', (3400, 3437), False, 'import torch\n'), ((5037, 5073), 'torch.nn.init.xavier_uniform_', 'nn.init.xavier_uniform_', (['params_data'], {}), '(params_data)\n', (5060, 5073), True, 'import torch.nn as nn\n'), ((5117, 5150), 'torch.nn.init.constant_', 'nn.init.constant_', (['params_data', '(0)'], {}), '(params_data, 0)\n', (5134, 5150), True, 'import torch.nn as nn\n'), ((5325, 5361), 'torch.nn.init.xavier_uniform_', 'nn.init.xavier_uniform_', (['params_data'], {}), '(params_data)\n', (5348, 5361), True, 'import torch.nn as nn\n'), ((5405, 5438), 'torch.nn.init.constant_', 'nn.init.constant_', (['params_data', '(0)'], {}), '(params_data, 0)\n', (5422, 5438), True, 'import torch.nn as nn\n'), ((5520, 5556), 'torch.nn.init.xavier_uniform_', 'nn.init.xavier_uniform_', (['params_data'], {}), '(params_data)\n', (5543, 5556), True, 'import torch.nn as nn\n'), ((5600, 5633), 'torch.nn.init.constant_', 'nn.init.constant_', (['params_data', '(0)'], {}), '(params_data, 0)\n', (5617, 5633), True, 'import torch.nn as nn\n'), ((7658, 7692), 'Coral.CORAL_torch', 'Coral.CORAL_torch', (['train_x', 'test_x'], {}), '(train_x, test_x)\n', (7675, 7692), False, 'import Coral\n'), ((10969, 10984), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (10982, 10984), False, 'import torch\n'), ((11387, 11402), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (11400, 11402), False, 'import torch\n'), ((11470, 11490), 'torch.max', 'torch.max', (['output', '(1)'], {}), '(output, 1)\n', (11479, 11490), False, 'import torch\n'), ((11885, 11947), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.PARAMS_FILE'], {}), '(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE)\n', (11897, 11947), False, 'import os\n'), ((12298, 12360), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.PARAMS_FILE'], {}), '(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE)\n', (12310, 12360), False, 'import os\n'), ((1747, 1808), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.NET1_ADABN'], {}), '(self.params_dir, cnnblstm_with_adabn.NET1_ADABN)\n', (1759, 1808), False, 'import os\n'), ((2437, 2498), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.NET2_ADABN'], {}), '(self.params_dir, cnnblstm_with_adabn.NET2_ADABN)\n', (2449, 2498), False, 'import os\n'), ((2726, 2787), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.NET3_ADABN'], {}), '(self.params_dir, cnnblstm_with_adabn.NET3_ADABN)\n', (2738, 2787), False, 'import os\n'), ((9730, 9750), 'torch.max', 'torch.max', (['output', '(1)'], {}), '(output, 1)\n', (9739, 9750), False, 'import torch\n'), ((3170, 3218), 'torch.zeros', 'torch.zeros', (['n_layers', 'batch_size', 'self.n_hidden'], {}), '(n_layers, batch_size, self.n_hidden)\n', (3181, 3218), False, 'import torch\n'), ((3242, 3290), 'torch.zeros', 'torch.zeros', (['n_layers', 'batch_size', 'self.n_hidden'], {}), '(n_layers, batch_size, self.n_hidden)\n', (3253, 3290), False, 'import torch\n'), ((11107, 11123), 'torch.autograd.Variable', 'Variable', (['test_x'], {}), '(test_x)\n', (11115, 11123), False, 'from torch.autograd import Variable\n'), ((11125, 11141), 'torch.autograd.Variable', 'Variable', (['test_y'], {}), '(test_y)\n', (11133, 11141), False, 'from torch.autograd import Variable\n'), ((8320, 8345), 'torch.cuda.device_count', 'torch.cuda.device_count', ([], {}), '()\n', (8343, 8345), False, 'import torch\n'), ((9262, 9275), 'torch.autograd.Variable', 'Variable', (['b_x'], {}), '(b_x)\n', (9270, 9275), False, 'from torch.autograd import Variable\n'), ((9277, 9290), 'torch.autograd.Variable', 'Variable', (['b_y'], {}), '(b_y)\n', (9285, 9290), False, 'from torch.autograd import Variable\n'), ((10875, 10900), 'torch.cuda.device_count', 'torch.cuda.device_count', ([], {}), '()\n', (10898, 10900), False, 'import torch\n'), ((12420, 12482), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.PARAMS_FILE'], {}), '(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE)\n', (12432, 12482), False, 'import os\n'), ((12567, 12629), 'os.path.join', 'os.path.join', (['self.params_dir', 'cnnblstm_with_adabn.PARAMS_FILE'], {}), '(self.params_dir, cnnblstm_with_adabn.PARAMS_FILE)\n', (12579, 12629), False, 'import os\n'), ((11028, 11044), 'torch.autograd.Variable', 'Variable', (['test_x'], {}), '(test_x)\n', (11036, 11044), False, 'from torch.autograd import Variable\n'), ((11053, 11069), 'torch.autograd.Variable', 'Variable', (['test_y'], {}), '(test_y)\n', (11061, 11069), False, 'from torch.autograd import Variable\n'), ((12499, 12519), 'torch.device', 'torch.device', (['"""cuda"""'], {}), "('cuda')\n", (12511, 12519), False, 'import torch\n'), ((12646, 12665), 'torch.device', 'torch.device', (['"""cpu"""'], {}), "('cpu')\n", (12658, 12665), False, 'import torch\n'), ((9193, 9206), 'torch.autograd.Variable', 'Variable', (['b_x'], {}), '(b_x)\n', (9201, 9206), False, 'from torch.autograd import Variable\n'), ((9215, 9228), 'torch.autograd.Variable', 'Variable', (['b_y'], {}), '(b_y)\n', (9223, 9228), False, 'from torch.autograd import Variable\n')]