nima_rafiee:

ResNet1D


@mrshenli   here is the code


import wget
import os
os.environ["CUDA_VISIBLE_DEVICES"]= "1,2,3,4,5,6"
import argparse
from tqdm import tqdm
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
import torch.distributed as dist
import torch.multiprocessing as mp
from apex import amp
from resent1D import ResNet1D
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel
import numpy as np
vis = True
try:
   from torch.utils.tensorboard import SummaryWriter
except:
   vis = False 
   
wt = 0.001




def train(rank, args):
   print("Passed GPU:" ,rank)
   dist.init_process_group(backend='nccl',init_method='env://',world_size=args.num_gpus ,rank=rank)    
   torch.cuda.set_device(rank)
   exp_num = f'h_{args.h_dim}_bs{args.bs}_hlr{args.h_lr}' 
   if vis == True :
       args.writer = SummaryWriter(f'runs/{exp_num}')  

   h_net = ResNet1D( in_channels=768, 
       base_filters=128, # 64 for ResNet1D, 352 for ResNeXt1D
       kernel_size=16, 
       stride=2, 
       groups=32, 
       n_block=48, 
       n_classes=args.h_dim, 
       downsample_gap=6, 
       increasefilter_gap=12, 
       use_do=True)

   h_net_opt = optim.Adam(h_net.parameters(), lr=args.h_lr, weight_decay=args.weight_decay)    
   
   h_net.cuda(rank) 
   model_lst , h_net_opt = amp.initialize([h_net,], h_net_opt, opt_level="O1")
   
   h_net = model_lst[0]
   h_net = DistributedDataParallel(h_net, device_ids=[rank], find_unused_parameters=True)    
   

   ds = TextDataset(dataset_name='AG_NEWS', out_cls=[])
   sampler = DistributedSampler(ds, num_replicas=args.num_gpus)
   loader = DataLoader(ds, batch_size=args.bs, sampler=sampler, shuffle=False)    
   
   
   loader = tqdm(loader)
   for epoch in range(args.epoch):
       sampler.set_epoch(epoch)        
       for i , (x1 , x2 , x_rnd, label) in enumerate(loader):
           h_net_opt.zero_grad()           
           h1 = h_net(x1)
           h2 = h_net(x2)
               
           loss = loss(h1, h2, temprature = args.temprature )             
           loss.backward()
           h_net_opt.step()
           loader.set_description(
               (
                   f' Epoch: {epoch + 1};  iter: {i} Loss: {loss.item()} '         
               )
           )

       if vis == True  :
           with torch.no_grad(): 
               args.writer.add_scalar("Loss", loss.item(), global_step=epoch, walltime=wt) 

       if rank == 0:
           #model_to_save = h_net.module if hasattr(h_net, 'module') else h_net  # Take care of distributed/parallel training
           torch.save(h_net.state_dict(), 'checkpoint/dd/h_net.checkpoint')
               
   args.writer.close()
   dist.destroy_process_group()
if __name__ == "__main__":
   
   parser = argparse.ArgumentParser()
   parser.add_argument('--bs', type=int, default=32)
   parser.add_argument('--epoch', type=int, default=50)    
   parser.add_argument('--h_lr', type=float, default=3e-4)   
   parser.add_argument('--h_dim', type=int, default=128)    
   parser.add_argument('--weight_decay', type=float, default=1e-6)    
   parser.add_argument('--temprature', type=float, default=0.5)        
      
   args = parser.parse_args()
   

   device_ids = list(range(torch.cuda.device_count()))
   gpus = len(device_ids)
   args.num_gpus = gpus 
       
   os.environ['MASTER_ADDR'] = '127.0.0.1'
   os.environ['MASTER_PORT'] = '8888'
   mp.spawn(train, nprocs=args.num_gpus, args=(args,))