I was trying to implement a simple CAE in Pytorch. I am seeing a lot of difference in train time in pytorch compared to tensorflow. Pytorch version is taking around 20 sec for 100 epochs whereas tensorflow version is taking around 5 sec for 100 epochs. Can anyone help me resolve this ? I have attached my code below.
Tensorflow code is similar to this https://github.com/panji1990/Deep-subspace-clustering-networks/blob/master/Pre-Train-Conv-AE-EYaleB.py
Also nvidia-smi shows 100% usage for pytorch version whereas it is around 45% for tensorflow version
class fullmodel(nn.Module):
def __init__(self,kernel_size,batch_size,num_class):
super(fullmodel, self).__init__()
self.batch_size = batch_size
self.enc1 = nn.Conv2d(1,10,kernel_size[0],padding = 2)
self.enc2 = nn.Conv2d(10,20,kernel_size[1],padding = 1)
self.enc3 = nn.Conv2d(20,30,kernel_size[2],padding = 1)
self.dec1 = nn.ConvTranspose2d(30,20,kernel_size = 3,padding = 1)
self.dec2 = nn.ConvTranspose2d(20,10,kernel_size = 3,padding = 1)
self.dec3 = nn.ConvTranspose2d(10,1,kernel_size = 5,padding = 2)
def forward(self,x):
enc_out = F.relu(self.enc3(F.relu(self.enc2(F.relu(self.enc1(x))))))
dec_out = F.relu(self.dec3(F.relu(self.dec2(F.relu(self.dec1(enc_out))))))
return enc_out,dec_out,
def train(model,data):
torch.backends.cudnn.benchmark = True
lr = 1e-3
num_epochs = 5000
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)
X = torch.from_numpy(data).float()
#data_train = torch.utils.data.DataLoader(dataset=X, batch_size=data.shape[0])
parameters = model.parameters()
model.apply(init_weights)
optim = torch.optim.Adam(model.parameters(), lr=lr)
model = model.to(device)
model = model.train()
X = X.to(device)
time1 = time.time()
for epoch in range(num_epochs):
avg_recon_loss = 0
latent,output = model(X)
recon_loss = 0.5*(torch.sum((output - X)**2))
optim.zero_grad()
recon_loss.backward()
optim.step()
avg_recon_loss += recon_loss
if (epoch+1)%100 == 0:
print("Iter : ",epoch+1)
print ("Loss : {:.4f}".format(avg_recon_loss/(38*64)))
print(time.time() - time1)
time1 = time.time()
if __name__ == "__main__":
data = sio.loadmat("../datasets/YaleBCrop025.mat")
imgs = data['Y']
I = []
Label = []
for i in range(imgs.shape[2]):
for j in range(imgs.shape[1]):
temp = np.reshape(imgs[:,j,i],[42,48])
Label.append(i)
I.append(temp)
I = np.array(I)
n_input = [42,48]
imgs = np.reshape(I, [I.shape[0], 1, n_input[0], n_input[1]])
kernel_size = [5,3,3]
classes = 38
batch_size = classes*64
model = fullmodel(kernel_size,classes,batch_size)
train(model,imgs)
