hc/JSCC
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

para modified

Browse Source
This commit is contained in:
chun 2023-12-23 14:32:47 +08:00
parent 34a6929650
commit 6d91d72581
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
train.py
View File

@ -82,6 +82,7 @@ def train(args: config_parser(), ratio: float, snr: float):
model = DeepJSCC(c=c, channel_type=args.channel, snr=snr) model = DeepJSCC(c=c, channel_type=args.channel, snr=snr)
if args.parallel and torch.cuda.device_count() > 1: if args.parallel and torch.cuda.device_count() > 1:
model = DataParallel(model, device_ids=list(range(torch.cuda.device_count()))) model = DataParallel(model, device_ids=list(range(torch.cuda.device_count())))
model.cuda()
criterion = nn.MSELoss(reduction='mean') criterion = nn.MSELoss(reduction='mean')
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
epoch_loop = tqdm(range(args.epochs), total=args.epochs, leave=False) epoch_loop = tqdm(range(args.epochs), total=args.epochs, leave=False)
@ -92,7 +93,6 @@ def train(args: config_parser(), ratio: float, snr: float):
optimizer.zero_grad() optimizer.zero_grad()
if not args.parallel: if not args.parallel:
images = images.cuda(device=device) images = images.cuda(device=device)
# images = images.cuda(device=device)
outputs = model(images) outputs = model(images)
loss = criterion(image_normalization('denormalization')(outputs), loss = criterion(image_normalization('denormalization')(outputs),
image_normalization('denormalization')(images)) image_normalization('denormalization')(images))